""" Shared LLM Utilities for BeatDebate Agents Consolidates LLM calling and JSON parsing patterns that are duplicated across agents, providing a single source of truth for LLM interactions. """ import json import re import asyncio from typing import Dict, Any, Optional, Union import structlog logger = structlog.get_logger(__name__) class LLMUtils: """ Shared utilities for LLM interactions across all agents. Consolidates: - LLM calling patterns - JSON response parsing - Error handling - Response cleaning and validation - Rate limiting for API quota management """ def __init__(self, llm_client, rate_limiter=None): """ Initialize LLM utilities with client and rate limiter. Args: llm_client: LLM client (e.g., Gemini client) rate_limiter: Rate limiter for API quota management """ self.llm_client = llm_client self.rate_limiter = rate_limiter self.logger = logger.bind(component="LLMUtils") async def call_llm_with_json_response( self, user_prompt: str, system_prompt: Optional[str] = None, max_retries: int = 2 ) -> Dict[str, Any]: """ Call LLM and parse JSON response with robust error handling and rate limiting. Args: user_prompt: User prompt for the LLM system_prompt: System prompt (optional) max_retries: Maximum number of retry attempts Returns: Parsed JSON response as dictionary Raises: ValueError: If JSON parsing fails after all retries RuntimeError: If LLM call fails """ for attempt in range(max_retries + 1): try: # Make LLM call with rate limiting response_text = await self._make_llm_call_with_rate_limiting(user_prompt, system_prompt) # Parse JSON response json_data = self._parse_json_response(response_text) self.logger.debug( "LLM JSON response parsed successfully", attempt=attempt + 1, response_keys=list(json_data.keys()) if isinstance(json_data, dict) else None ) return json_data except json.JSONDecodeError as e: self.logger.warning( "JSON parsing failed", attempt=attempt + 1, error=str(e), response_preview=response_text[:200] if 'response_text' in locals() else None ) if attempt == max_retries: # Try alternative parsing methods on final attempt try: return self._aggressive_json_parsing(response_text) except Exception: raise ValueError(f"Failed to parse JSON after {max_retries + 1} attempts: {e}") except Exception as e: error_msg = str(e) if "429" in error_msg or "quota" in error_msg.lower() or "rate limit" in error_msg.lower(): # Handle rate limiting with exponential backoff wait_time = min(60, 2 ** attempt) # Cap at 60 seconds self.logger.warning( "Rate limit hit, waiting before retry", attempt=attempt + 1, wait_time=wait_time, error=error_msg ) if attempt < max_retries: await asyncio.sleep(wait_time) continue self.logger.error( "LLM call failed", attempt=attempt + 1, error=error_msg ) if attempt == max_retries: raise RuntimeError(f"LLM call failed after {max_retries + 1} attempts: {e}") # This should never be reached, but just in case raise RuntimeError("Unexpected error in LLM call loop") async def call_llm( self, user_prompt: str, system_prompt: Optional[str] = None ) -> str: """ Call LLM and return raw text response with rate limiting. Args: user_prompt: User prompt for the LLM system_prompt: System prompt (optional) Returns: Raw LLM response text Raises: RuntimeError: If LLM call fails """ try: response_text = await self._make_llm_call_with_rate_limiting(user_prompt, system_prompt) self.logger.debug( "LLM text response received", response_length=len(response_text) ) return response_text except Exception as e: self.logger.error("LLM call failed", error=str(e)) raise RuntimeError(f"LLM call failed: {e}") async def _make_llm_call_with_rate_limiting( self, user_prompt: str, system_prompt: Optional[str] = None ) -> str: """ Make LLM call with rate limiting and unified error handling. Args: user_prompt: User prompt system_prompt: System prompt Returns: LLM response text """ if not self.llm_client: raise RuntimeError("LLM client not initialized") # Apply rate limiting if available if self.rate_limiter: await self.rate_limiter.wait_if_needed() try: # Combine system and user prompts full_prompt = ( f"{system_prompt}\n\n{user_prompt}" if system_prompt else user_prompt ) self.logger.debug( "Making LLM call", prompt_length=len(full_prompt), has_system_prompt=system_prompt is not None ) # Call LLM - handle both sync and async clients response = self.llm_client.generate_content(full_prompt) # If it's a coroutine (async client), await it if hasattr(response, '__await__'): response = await response return response.text except Exception as e: self.logger.error("LLM API call failed", error=str(e)) raise def _parse_json_response(self, response_text: str) -> Dict[str, Any]: """ Parse JSON response with robust error handling and cleaning. Args: response_text: Raw LLM response text Returns: Parsed JSON data Raises: json.JSONDecodeError: If JSON parsing fails """ try: # Clean the response text cleaned_text = self._clean_response_text(response_text) # Extract JSON boundaries json_str = self._extract_json_boundaries(cleaned_text) # Additional JSON cleaning for common LLM issues json_str = self._clean_json_string(json_str) # Parse JSON json_data = json.loads(json_str) self.logger.debug( "JSON parsing successful", original_length=len(response_text), cleaned_length=len(json_str), keys=list(json_data.keys()) if isinstance(json_data, dict) else None ) return json_data except json.JSONDecodeError as e: self.logger.warning( "Initial JSON parsing failed", error=str(e), response_preview=response_text[:300] ) raise def _clean_response_text(self, response_text: str) -> str: """Clean response text by removing markdown and explanatory text.""" cleaned = response_text.strip() # Remove markdown code blocks if cleaned.startswith('```'): lines = cleaned.split('\n') # Remove first line if it's markdown if lines[0].startswith('```'): lines = lines[1:] # Remove last line if it's markdown if lines and lines[-1].startswith('```'): lines = lines[:-1] cleaned = '\n'.join(lines) return cleaned.strip() def _extract_json_boundaries(self, text: str) -> str: """Extract JSON object boundaries from text.""" # Find the first opening brace start_idx = text.find('{') if start_idx == -1: raise ValueError("No JSON object found in response") # Find matching closing brace by counting braces brace_count = 0 end_idx = start_idx for i, char in enumerate(text[start_idx:], start_idx): if char == '{': brace_count += 1 elif char == '}': brace_count -= 1 if brace_count == 0: end_idx = i + 1 break if brace_count != 0: # If braces don't match, try to find the last closing brace end_idx = text.rfind('}') if end_idx == -1 or end_idx <= start_idx: raise ValueError("Unmatched braces in JSON response") end_idx += 1 return text[start_idx:end_idx] def _clean_json_string(self, json_str: str) -> str: """Clean JSON string to fix common LLM formatting issues.""" # Remove any trailing commas before closing braces/brackets json_str = re.sub(r',(\s*[}\]])', r'\1', json_str) # Remove any comments (// or /* */) json_str = re.sub(r'//.*$', '', json_str, flags=re.MULTILINE) json_str = re.sub(r'/\*.*?\*/', '', json_str, flags=re.DOTALL) # Fix common typos in boolean/null values json_str = re.sub(r'\btrue\b', 'true', json_str, flags=re.IGNORECASE) json_str = re.sub(r'\bfalse\b', 'false', json_str, flags=re.IGNORECASE) json_str = re.sub(r'\bnull\b', 'null', json_str, flags=re.IGNORECASE) # Replace single quotes with double quotes for keys and string values # This is a simple approach - for complex cases, we'd need a proper parser json_str = re.sub(r"'([^']*)':", r'"\1":', json_str) # Keys json_str = re.sub(r":\s*'([^']*)'", r': "\1"', json_str) # String values return json_str def _aggressive_json_parsing(self, response_text: str) -> Dict[str, Any]: """ Aggressive JSON parsing as a last resort. Args: response_text: Raw response text Returns: Parsed JSON data Raises: ValueError: If all parsing attempts fail """ self.logger.info("Attempting aggressive JSON parsing") # Attempt 1: Try fixing common JSON issues try: fixed_json = self._fix_common_json_issues(response_text) return json.loads(fixed_json) except Exception as e: self.logger.debug("Fixed JSON parsing failed", error=str(e)) # Attempt 2: Use regex to extract JSON-like structure try: extracted_json = self._extract_json_with_regex(response_text) if extracted_json: return json.loads(extracted_json) except Exception as e: self.logger.debug("Regex JSON extraction failed", error=str(e)) # Attempt 3: Try to build JSON from key-value patterns try: constructed_json = self._construct_json_from_patterns(response_text) if constructed_json: return constructed_json except Exception as e: self.logger.debug("Pattern-based JSON construction failed", error=str(e)) raise ValueError("All aggressive JSON parsing attempts failed") def _fix_common_json_issues(self, response_text: str) -> str: """Attempt to fix common JSON formatting issues.""" # Find JSON boundaries more aggressively start_idx = response_text.find('{') if start_idx == -1: return response_text # Extract everything from first { to last } end_idx = response_text.rfind('}') if end_idx == -1: return response_text json_candidate = response_text[start_idx:end_idx + 1] # Apply aggressive cleaning json_candidate = self._clean_json_string(json_candidate) # Remove any text before first { or after last } json_candidate = re.sub(r'^[^{]*', '', json_candidate) json_candidate = re.sub(r'}[^}]*$', '}', json_candidate) return json_candidate def _extract_json_with_regex(self, response_text: str) -> Optional[str]: """Extract JSON using regex patterns as a last resort.""" # Look for JSON-like structure with balanced braces pattern = r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}' matches = re.findall(pattern, response_text, re.DOTALL) if matches: # Return the longest match (most likely to be complete) longest_match = max(matches, key=len) return self._clean_json_string(longest_match) return None def _construct_json_from_patterns(self, response_text: str) -> Optional[Dict[str, Any]]: """Construct JSON from key-value patterns in text.""" try: # Look for key-value patterns like "key": "value" or "key": value kv_pattern = r'"([^"]+)":\s*(?:"([^"]*)"|([^,}\s]+))' matches = re.findall(kv_pattern, response_text) if matches: result = {} for key, str_value, other_value in matches: value = str_value if str_value else other_value # Try to convert to appropriate type if value.lower() == 'true': result[key] = True elif value.lower() == 'false': result[key] = False elif value.lower() == 'null': result[key] = None elif value.isdigit(): result[key] = int(value) elif self._is_float(value): result[key] = float(value) else: result[key] = value return result if result else None except Exception as e: self.logger.debug("Pattern-based JSON construction failed", error=str(e)) return None def _is_float(self, value: str) -> bool: """Check if string represents a float.""" try: float(value) return True except ValueError: return False def validate_json_structure( self, json_data: Dict[str, Any], required_keys: Optional[list] = None, optional_keys: Optional[list] = None ) -> Dict[str, Any]: """ Validate and enhance JSON structure. Args: json_data: Parsed JSON data required_keys: List of required keys optional_keys: List of optional keys to set defaults for Returns: Validated and enhanced JSON data """ if not isinstance(json_data, dict): raise ValueError("JSON data must be a dictionary") # Check required keys if required_keys: missing_keys = [key for key in required_keys if key not in json_data] if missing_keys: self.logger.warning("Missing required keys", missing_keys=missing_keys) # Set default values for missing required keys for key in missing_keys: json_data[key] = self._get_default_value_for_key(key) # Set defaults for optional keys if optional_keys: for key in optional_keys: if key not in json_data: json_data[key] = self._get_default_value_for_key(key) self.logger.debug( "JSON structure validated", keys=list(json_data.keys()), required_keys=required_keys, optional_keys=optional_keys ) return json_data def _get_default_value_for_key(self, key: str) -> Union[str, list, dict, int, float]: """Get appropriate default value based on key name.""" # Common key patterns and their default values if 'list' in key.lower() or key.endswith('s'): return [] elif 'dict' in key.lower() or 'entities' in key.lower(): return {} elif 'count' in key.lower() or 'score' in key.lower(): return 0 elif 'confidence' in key.lower(): return 0.0 elif 'intent' in key.lower(): return 'discovery' elif 'complexity' in key.lower(): return 'medium' else: return "" def create_structured_prompt( self, task_description: str, input_data: Dict[str, Any], output_format: Dict[str, Any], examples: Optional[list] = None ) -> str: """ Create a structured prompt for LLM with consistent formatting. Args: task_description: Description of the task input_data: Input data to include in prompt output_format: Expected output format examples: Optional examples to include Returns: Formatted prompt string """ prompt_parts = [ f"Task: {task_description}", "", "Input Data:", json.dumps(input_data, indent=2), "", "Required Output Format:", json.dumps(output_format, indent=2) ] if examples: prompt_parts.extend([ "", "Examples:", json.dumps(examples, indent=2) ]) prompt_parts.extend([ "", "Please provide your response in the exact JSON format specified above.", "Ensure all required fields are included and properly formatted." ]) return "\n".join(prompt_parts) async def generate_reasoning( self, reasoning_prompt: str, max_tokens: int = 100, temperature: float = 0.7 ) -> str: """ Generate reasoning text using LLM. Args: reasoning_prompt: Prompt for reasoning generation max_tokens: Maximum tokens to generate temperature: Temperature for generation Returns: Generated reasoning text Raises: RuntimeError: If LLM call fails """ try: response_text = await self.call_llm(reasoning_prompt) # Clean and truncate response if needed cleaned_response = response_text.strip() # Basic truncation if too long (rough token estimation) if len(cleaned_response) > max_tokens * 4: # Rough estimate: 4 chars per token cleaned_response = cleaned_response[:max_tokens * 4] + "..." self.logger.debug( "Reasoning generated successfully", response_length=len(cleaned_response) ) return cleaned_response except Exception as e: self.logger.error("Reasoning generation failed", error=str(e)) raise RuntimeError(f"Failed to generate reasoning: {e}")