""" DSPy Streaming Adapter for Snakepit This adapter provides streaming implementations of DSPy operations, enabling progressive result delivery from Python to Elixir. """ import json import time import traceback from typing import Iterator, Dict, Any, List, Optional import logging logger = logging.getLogger(__name__) # Import DSPy when available try: import dspy DSPY_AVAILABLE = True except ImportError: DSPY_AVAILABLE = False logger.warning("DSPy not available. Streaming commands will use mock implementations.") class DSPyStreamingHandler: """Handler for streaming DSPy operations""" def __init__(self): self.supported_streaming_commands = [ "stream_predict_batch", "stream_chain_of_thought", "stream_react_steps", "stream_optimization_progress", "stream_retrieval_results" ] # Store any configured LM self.lm = None if DSPY_AVAILABLE: # Check if an LM is already configured try: self.lm = dspy.settings.lm except: pass def supports_streaming(self) -> bool: """Check if streaming is supported""" return True def get_streaming_commands(self) -> List[str]: """Return list of supported streaming commands""" return self.supported_streaming_commands def process_stream_command(self, command: str, args: Dict[str, Any]) -> Iterator[Dict[str, Any]]: """Main entry point for streaming commands""" # Ensure we have a language model configured if DSPY_AVAILABLE and not self.lm: # Try to get LM from stored objects if hasattr(self, 'protocol_handler') and hasattr(self.protocol_handler, 'stored_objects'): stored_lm = self.protocol_handler.stored_objects.get('default_lm') if stored_lm: dspy.configure(lm=stored_lm) self.lm = stored_lm try: if command == "stream_predict_batch": yield from self._stream_predict_batch(args) elif command == "stream_chain_of_thought": yield from self._stream_chain_of_thought(args) elif command == "stream_react_steps": yield from self._stream_react_steps(args) elif command == "stream_optimization_progress": yield from self._stream_optimization_progress(args) elif command == "stream_retrieval_results": yield from self._stream_retrieval_results(args) else: yield {"error": f"Unknown streaming command: {command}"} except Exception as e: yield { "error": str(e), "traceback": traceback.format_exc(), "type": "stream_error" } def _stream_predict_batch(self, args: Dict[str, Any]) -> Iterator[Dict[str, Any]]: """Stream predictions for a batch of inputs""" signature = args.get('signature', 'input -> output') items = args.get('items', []) batch_size = args.get('batch_size', 1) if not DSPY_AVAILABLE: # Mock implementation for i, item in enumerate(items): yield { 'type': 'prediction', 'index': i, 'total': len(items), 'input': item, 'output': {'output': f"Mock prediction for: {item.get('input', str(item))}"}, 'progress': (i + 1) / len(items) } time.sleep(0.1) # Simulate processing time yield {'type': 'complete', 'total': len(items)} return # Real DSPy implementation predictor = dspy.Predict(signature) # Process in batches for batch_start in range(0, len(items), batch_size): batch_end = min(batch_start + batch_size, len(items)) batch = items[batch_start:batch_end] for i, item in enumerate(batch): global_index = batch_start + i try: # Execute prediction prediction = predictor(**item) # Convert to dict output = prediction.toDict() if hasattr(prediction, 'toDict') else { k: v for k, v in prediction.__dict__.items() if not k.startswith('_') } yield { 'type': 'prediction', 'index': global_index, 'total': len(items), 'input': item, 'output': output, 'progress': (global_index + 1) / len(items), 'batch': batch_start // batch_size + 1 } except Exception as e: yield { 'type': 'error', 'index': global_index, 'error': str(e), 'input': item } yield {'type': 'complete', 'total': len(items)} def _stream_chain_of_thought(self, args: Dict[str, Any]) -> Iterator[Dict[str, Any]]: """Stream reasoning steps from Chain of Thought""" signature = args.get('signature', 'question -> answer') question = args.get('question', '') max_reasoning_steps = args.get('max_reasoning_steps', 10) if not DSPY_AVAILABLE: # Mock implementation with simulated reasoning yield {'type': 'thinking', 'status': 'Analyzing question...'} time.sleep(0.5) mock_steps = [ "Breaking down the problem...", "Identifying key components...", "Applying logical reasoning...", "Formulating answer..." ] for i, step in enumerate(mock_steps): yield { 'type': 'reasoning_step', 'step': i + 1, 'content': step } time.sleep(0.3) yield { 'type': 'final_answer', 'answer': f"Mock answer for: {question}", 'full_reasoning': '\n'.join(mock_steps) } return # Real DSPy implementation yield {'type': 'thinking', 'status': 'Initializing Chain of Thought...'} cot = dspy.ChainOfThought(signature) # Hook into internal reasoning process if possible reasoning_buffer = [] def capture_reasoning(text): """Capture reasoning steps as they're generated""" reasoning_buffer.append(text) yield { 'type': 'reasoning_fragment', 'content': text, 'step': len(reasoning_buffer) } yield {'type': 'thinking', 'status': 'Generating reasoning...'} # Execute with potential hooks try: result = cot(question=question) # Process reasoning if available if hasattr(result, 'reasoning') and result.reasoning: reasoning_lines = result.reasoning.split('\n') for i, line in enumerate(reasoning_lines): if line.strip(): yield { 'type': 'reasoning_step', 'step': i + 1, 'content': line.strip() } # Final answer answer = result.answer if hasattr(result, 'answer') else str(result) yield { 'type': 'final_answer', 'answer': answer, 'full_reasoning': result.reasoning if hasattr(result, 'reasoning') else None, 'completions': result.completions if hasattr(result, 'completions') else None } except Exception as e: yield { 'type': 'error', 'error': str(e), 'partial_reasoning': '\n'.join(reasoning_buffer) if reasoning_buffer else None } def _stream_react_steps(self, args: Dict[str, Any]) -> Iterator[Dict[str, Any]]: """Stream ReAct reasoning and tool usage steps""" signature = args.get('signature', 'question -> answer') question = args.get('question', '') tools = args.get('tools', []) max_iters = args.get('max_iters', 5) if not DSPY_AVAILABLE: # Mock implementation yield {'type': 'init', 'status': 'Initializing ReAct...'} mock_steps = [ {'action': 'think', 'thought': 'I need to analyze this question'}, {'action': 'tool', 'tool': 'search', 'input': question[:20] + '...'}, {'action': 'observe', 'observation': 'Found relevant information'}, {'action': 'think', 'thought': 'Based on the results, I can conclude...'} ] for i, step in enumerate(mock_steps): yield { 'type': 'react_step', 'iteration': i + 1, 'action': step['action'], 'content': step.get('thought') or step.get('tool') or step.get('observation', '') } time.sleep(0.4) yield { 'type': 'final_answer', 'answer': f"Mock ReAct answer for: {question}", 'iterations': len(mock_steps) } return # Real DSPy implementation would go here # For now, return a placeholder yield { 'type': 'error', 'error': 'ReAct streaming not yet implemented with real DSPy' } def _stream_optimization_progress(self, args: Dict[str, Any]) -> Iterator[Dict[str, Any]]: """Stream optimization progress updates""" optimizer_type = args.get('optimizer', 'BootstrapFewShot') dataset = args.get('dataset', []) metric = args.get('metric', 'accuracy') if not DSPY_AVAILABLE: # Mock optimization progress total_steps = 10 for i in range(total_steps): yield { 'type': 'optimization_step', 'step': i + 1, 'total_steps': total_steps, 'metric_value': 0.5 + (0.4 * i / total_steps), # Simulate improvement 'best_score': 0.5 + (0.4 * i / total_steps), 'status': f'Optimizing... ({i+1}/{total_steps})' } time.sleep(0.5) yield { 'type': 'optimization_complete', 'final_score': 0.9, 'total_steps': total_steps, 'best_prompts': ['Optimized prompt 1', 'Optimized prompt 2'] } return # Real optimization would hook into DSPy optimizers yield { 'type': 'error', 'error': 'Optimization streaming not yet implemented with real DSPy' } def _stream_retrieval_results(self, args: Dict[str, Any]) -> Iterator[Dict[str, Any]]: """Stream retrieval results as they're found""" query = args.get('query', '') k = args.get('k', 10) retriever_type = args.get('retriever', 'ColBERTv2') # Mock retrieval results for i in range(k): yield { 'type': 'retrieval_result', 'index': i, 'document': { 'id': f'doc_{i}', 'content': f'Document {i} content related to: {query[:30]}...', 'score': 0.9 - (i * 0.05) }, 'progress': (i + 1) / k } time.sleep(0.1) yield { 'type': 'retrieval_complete', 'total_results': k, 'query': query } def create_handler(): """Factory function to create the streaming handler""" return DSPyStreamingHandler()