kollabor 0.4.9__py3-none-any.whl

core/llm/response_processor.py

@@ -0,0 +1,364 @@
+"""Response processing and formatting for LLM outputs.
+
+Handles response parsing, formatting, and special tag processing
+for LLM responses including thinking tags and final responses.
+"""
+
+import json
+import logging
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+class ResponseProcessor:
+    """Process and format LLM responses.
+    
+    Handles parsing of special tags, response formatting,
+    and content extraction from raw LLM outputs.
+    """
+    
+    def __init__(self):
+        """Initialize response processor."""
+        # Regex patterns for tag extraction
+        self.thinking_pattern = re.compile(
+            r'<think>(.*?)</think>', 
+            re.DOTALL | re.IGNORECASE
+        )
+        self.terminal_pattern = re.compile(
+            r'<terminal>(.*?)</terminal>', 
+            re.DOTALL | re.IGNORECASE
+        )
+        
+        logger.info("Response processor initialized")
+    
+    def process_response(self, raw_response: str) -> Dict[str, Any]:
+        """Process raw LLM response.
+        
+        Args:
+            raw_response: Raw response from LLM
+            
+        Returns:
+            Processed response with extracted components
+        """
+        # Extract components
+        thinking_content = self._extract_thinking(raw_response)
+        terminal_commands = self._extract_terminal_commands(raw_response)
+        
+        # Clean main content
+        clean_content = self._clean_content(raw_response)
+        
+        # Determine response type
+        response_type = self._determine_response_type(
+            raw_response, 
+            thinking_content, 
+            terminal_commands
+        )
+        
+        processed = {
+            "type": response_type,
+            "content": clean_content,
+            "raw": raw_response,
+            "components": {
+                "thinking": thinking_content,
+                "terminal_commands": terminal_commands
+            },
+            "metadata": {
+                "has_thinking": bool(thinking_content),
+                "has_terminal_commands": bool(terminal_commands),
+                "word_count": len(clean_content.split()),
+                "char_count": len(clean_content)
+            }
+        }
+        
+        logger.debug(f"Processed response: type={response_type}, components={len(processed['components'])}")
+        return processed
+    
+    def _extract_thinking(self, content: str) -> List[str]:
+        """Extract thinking tags from content.
+        
+        Args:
+            content: Raw content with potential thinking tags
+            
+        Returns:
+            List of thinking content blocks
+        """
+        matches = self.thinking_pattern.findall(content)
+        return [match.strip() for match in matches if match.strip()]
+    
+    def _extract_terminal_commands(self, content: str) -> List[str]:
+        """Extract terminal commands from content.
+        
+        Args:
+            content: Raw content with potential terminal tags
+            
+        Returns:
+            List of terminal commands
+        """
+        matches = self.terminal_pattern.findall(content)
+        return [match.strip() for match in matches if match.strip()]
+    
+    def _extract_tool_calls(self, content: str) -> List[Dict[str, Any]]:
+        """Extract tool calls from content.
+        
+        Args:
+            content: Raw content with potential tool calls
+            
+        Returns:
+            List of parsed tool calls
+        """
+        tool_calls = []
+        matches = self.tool_call_pattern.findall(content)
+        
+        for match in matches:
+            try:
+                # Try to parse as JSON
+                tool_data = json.loads(match)
+                tool_calls.append(tool_data)
+            except json.JSONDecodeError:
+                # Parse as plain text command
+                lines = match.strip().split('\n')
+                if lines:
+                    tool_calls.append({
+                        "command": lines[0],
+                        "raw": match
+                    })
+        
+        return tool_calls
+    
+    def _extract_code_blocks(self, content: str) -> List[Dict[str, str]]:
+        """Extract code blocks from content.
+        
+        Args:
+            content: Raw content with potential code blocks
+            
+        Returns:
+            List of code blocks with language and content
+        """
+        code_blocks = []
+        matches = self.code_block_pattern.findall(content)
+        
+        for language, code in matches:
+            code_blocks.append({
+                "language": language or "plain",
+                "code": code.strip()
+            })
+        
+        return code_blocks
+    
+    def _clean_content(self, content: str) -> str:
+        """Remove special tags from content.
+        
+        Args:
+            content: Raw content with tags
+            
+        Returns:
+            Cleaned content without special tags
+        """
+        # Remove thinking tags
+        cleaned = self.thinking_pattern.sub('', content)
+        
+        # Remove final response tags but keep content
+        final_match = self.final_response_pattern.search(cleaned)
+        if final_match:
+            cleaned = self.final_response_pattern.sub(final_match.group(1), cleaned)
+        
+        # Remove tool call tags
+        cleaned = self.tool_call_pattern.sub('', cleaned)
+        
+        # Clean up extra whitespace
+        cleaned = re.sub(r'\n{3,}', '\n\n', cleaned)
+        cleaned = cleaned.strip()
+        
+        return cleaned
+    
+    def _determine_response_type(self, raw: str, thinking: List[str], 
+                                 final: Optional[str], tools: List[Dict]) -> str:
+        """Determine the type of response.
+        
+        Args:
+            raw: Raw response content
+            thinking: Extracted thinking content
+            final: Extracted final response
+            tools: Extracted tool calls
+            
+        Returns:
+            Response type identifier
+        """
+        if final:
+            return "final"
+        elif tools:
+            return "tool_execution"
+        elif thinking:
+            return "reasoning"
+        elif "?" in raw:
+            return "question"
+        elif len(raw) < 100:
+            return "brief"
+        else:
+            return "standard"
+    
+    def format_for_display(self, processed_response: Dict[str, Any]) -> str:
+        """Format processed response for display.
+        
+        Args:
+            processed_response: Processed response data
+            
+        Returns:
+            Formatted string for display
+        """
+        components = processed_response.get("components", {})
+        
+        # Build display string
+        display_parts = []
+        
+        # Add thinking if present (dimmed)
+        thinking = components.get("thinking", [])
+        if thinking:
+            for thought in thinking:
+                display_parts.append(f"[dim]{thought}[/dim]")
+                display_parts.append("")
+        
+        # Add main content or final response
+        if components.get("final_response"):
+            display_parts.append(components["final_response"])
+        else:
+            display_parts.append(processed_response.get("content", ""))
+        
+        return "\n".join(display_parts)
+    
+    def extract_metrics(self, processed_response: Dict[str, Any]) -> Dict[str, Any]:
+        """Extract metrics from processed response.
+        
+        Args:
+            processed_response: Processed response data
+            
+        Returns:
+            Response metrics
+        """
+        metadata = processed_response.get("metadata", {})
+        components = processed_response.get("components", {})
+        
+        metrics = {
+            "response_type": processed_response.get("type", "unknown"),
+            "word_count": metadata.get("word_count", 0),
+            "char_count": metadata.get("char_count", 0),
+            "thinking_blocks": len(components.get("thinking", [])),
+            "code_blocks": len(components.get("code_blocks", [])),
+            "tool_calls": len(components.get("tool_calls", [])),
+            "has_final_response": metadata.get("has_final_response", False),
+            "complexity": self._assess_complexity(processed_response)
+        }
+        
+        return metrics
+    
+    def _assess_complexity(self, processed_response: Dict[str, Any]) -> str:
+        """Assess response complexity.
+        
+        Args:
+            processed_response: Processed response data
+            
+        Returns:
+            Complexity level (simple, moderate, complex)
+        """
+        metadata = processed_response.get("metadata", {})
+        
+        # Calculate complexity score
+        score = 0
+        
+        if metadata.get("word_count", 0) > 200:
+            score += 2
+        elif metadata.get("word_count", 0) > 100:
+            score += 1
+        
+        if metadata.get("has_thinking"):
+            score += 2
+        
+        if metadata.get("has_code"):
+            score += 1
+        
+        if metadata.get("has_tool_calls"):
+            score += 1
+        
+        # Map score to complexity level
+        if score >= 4:
+            return "complex"
+        elif score >= 2:
+            return "moderate"
+        else:
+            return "simple"
+    
+    def merge_streaming_chunks(self, chunks: List[str]) -> str:
+        """Merge streaming response chunks.
+        
+        Args:
+            chunks: List of response chunks
+            
+        Returns:
+            Merged response content
+        """
+        # Join chunks
+        merged = "".join(chunks)
+        
+        # Fix any broken tags from chunking
+        merged = self._fix_broken_tags(merged)
+        
+        return merged
+    
+    def _fix_broken_tags(self, content: str) -> str:
+        """Fix broken tags from streaming.
+        
+        Args:
+            content: Content with potentially broken tags
+            
+        Returns:
+            Content with fixed tags
+        """
+        # Fix broken thinking tags
+        content = re.sub(r'<thi\s*nk>', '<think>', content, flags=re.IGNORECASE)
+        content = re.sub(r'</thi\s*nk>', '</think>', content, flags=re.IGNORECASE)
+        
+        # Fix broken final response tags
+        content = re.sub(r'<final_res\s*ponse>', '<final_response>', content, flags=re.IGNORECASE)
+        content = re.sub(r'</final_res\s*ponse>', '</final_response>', content, flags=re.IGNORECASE)
+        
+        return content
+    
+    def validate_response(self, response: str) -> Tuple[bool, List[str]]:
+        """Validate response format and content.
+        
+        Args:
+            response: Response to validate
+            
+        Returns:
+            Tuple of (is_valid, list_of_issues)
+        """
+        issues = []
+        
+        # Check for unclosed tags
+        if '<think>' in response and '</think>' not in response:
+            issues.append("Unclosed thinking tag")
+        
+        if '<final_response>' in response and '</final_response>' not in response:
+            issues.append("Unclosed final response tag")
+        
+        # Check for empty response
+        if not response.strip():
+            issues.append("Empty response")
+        
+        # Check for truncation indicators
+        if response.endswith('...') and len(response) > 1000:
+            issues.append("Response appears truncated")
+        
+        # Check for malformed JSON in tool calls
+        tool_matches = self.tool_call_pattern.findall(response)
+        for match in tool_matches:
+            if match.strip().startswith('{'):
+                try:
+                    json.loads(match)
+                except json.JSONDecodeError:
+                    issues.append("Malformed JSON in tool call")
+        
+        is_valid = len(issues) == 0
+        return is_valid, issues