#!/usr/bin/env python3 """ DSPy general purpose bridge for Snakepit. Handles basic DSPy operations like Predict, ChainOfThought, etc. """ import json import sys import traceback import logging from datetime import datetime try: import dspy except ImportError: print("Error: dspy-ai package not installed. Run: pip install dspy-ai", file=sys.stderr) sys.exit(1) # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) class DSPyBridge: """Bridge between Elixir and DSPy for general operations.""" def __init__(self): self.modules = {} self.signatures = {} self.lm = None def initialize_lm(self, provider="openai", **kwargs): """Initialize the language model.""" if provider == "openai": self.lm = dspy.OpenAI(**kwargs) elif provider == "anthropic": self.lm = dspy.Claude(**kwargs) else: raise ValueError(f"Unknown provider: {provider}") dspy.settings.configure(lm=self.lm) logger.info(f"Initialized LM provider: {provider}") def handle_request(self, request): """Handle incoming request from Elixir.""" try: operation = request.get("operation") args = request.get("args", {}) opts = request.get("opts", {}) logger.debug(f"Handling operation: {operation}") # Route to appropriate handler if operation == "dspy.Predict": return self.handle_predict(args, opts) elif operation == "dspy.ChainOfThought": return self.handle_chain_of_thought(args, opts) elif operation == "dspy.ReAct": return self.handle_react(args, opts) elif operation == "dspy.ProgramOfThought": return self.handle_program_of_thought(args, opts) elif operation == "initialize_lm": return self.handle_initialize_lm(args) else: return self.error_response(f"Unknown operation: {operation}") except Exception as e: logger.error(f"Error handling request: {str(e)}", exc_info=True) return self.error_response(str(e), traceback.format_exc()) def handle_predict(self, args, opts): """Handle basic prediction.""" signature = self.get_or_create_signature(args["signature"]) inputs = args["inputs"] # Create predictor predictor = dspy.Predict(signature) # Execute prediction result = predictor(**inputs) # Extract outputs output = {} for field in signature.output_fields: if hasattr(result, field): output[field] = getattr(result, field) return self.success_response(output) def handle_chain_of_thought(self, args, opts): """Handle chain of thought reasoning.""" signature = self.get_or_create_signature(args["signature"]) inputs = args["inputs"] # Create CoT module cot = dspy.ChainOfThought(signature) # Execute result = cot(**inputs) # Extract outputs including reasoning output = {} for field in signature.output_fields: if hasattr(result, field): output[field] = getattr(result, field) # Include reasoning if available if hasattr(result, "rationale"): output["reasoning"] = result.rationale return self.success_response(output) def handle_react(self, args, opts): """Handle ReAct pattern.""" # This is a placeholder - full ReAct implementation needs tool integration return self.error_response("ReAct not yet implemented") def handle_program_of_thought(self, args, opts): """Handle program of thought.""" # This is a placeholder return self.error_response("ProgramOfThought not yet implemented") def handle_initialize_lm(self, args): """Initialize language model configuration.""" provider = args.get("provider", "openai") config = args.get("config", {}) self.initialize_lm(provider, **config) return self.success_response({"status": "initialized", "provider": provider}) def get_or_create_signature(self, sig_spec): """Get or create a DSPy signature.""" if isinstance(sig_spec, str): # Simple string signature if sig_spec not in self.signatures: self.signatures[sig_spec] = dspy.Signature(sig_spec) return self.signatures[sig_spec] else: # Complex signature with fields sig_key = self._signature_key(sig_spec) if sig_key not in self.signatures: # Create signature from specification fields = [] # Add input fields for field in sig_spec.get("inputs", []): fields.append(dspy.InputField( prefix=field["name"], desc=field.get("description", "") )) # Add output fields for field in sig_spec.get("outputs", []): fields.append(dspy.OutputField( prefix=field["name"], desc=field.get("description", "") )) self.signatures[sig_key] = dspy.Signature(*fields) return self.signatures[sig_key] def _signature_key(self, sig_spec): """Generate a key for signature caching.""" return json.dumps(sig_spec, sort_keys=True) def success_response(self, result): """Create a success response.""" return { "success": True, "result": result, "timestamp": datetime.utcnow().isoformat() } def error_response(self, error, traceback=None): """Create an error response.""" response = { "success": False, "error": error, "timestamp": datetime.utcnow().isoformat() } if traceback: response["traceback"] = traceback return response def main(): """Main entry point for Snakepit worker.""" logger.info("DSPy bridge starting...") bridge = DSPyBridge() # Initialize default LM if environment variables are set if "OPENAI_API_KEY" in os.environ: bridge.initialize_lm("openai") # Main request/response loop while True: try: # Read request from stdin (Snakepit protocol) request = read_request() if request is None: break # Handle request response = bridge.handle_request(request) # Write response write_response(response) except KeyboardInterrupt: logger.info("Shutting down...") break except Exception as e: logger.error(f"Fatal error: {str(e)}", exc_info=True) write_response({ "success": False, "error": f"Fatal error: {str(e)}" }) def read_request(): """Read a request using Snakepit protocol.""" try: # Snakepit sends JSON lines line = sys.stdin.readline() if not line: return None return json.loads(line) except json.JSONDecodeError as e: logger.error(f"Invalid JSON: {e}") return None def write_response(response): """Write a response using Snakepit protocol.""" # Snakepit expects JSON lines json_str = json.dumps(response) sys.stdout.write(json_str + "\n") sys.stdout.flush() if __name__ == "__main__": import os main()