""" Base Agent Class All agents inherit from this to provide consistent interface """ import asyncio import json import time import logging from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Dict, Any, Optional from datetime import datetime logger = logging.getLogger(__name__) @dataclass class AgentConfig: """Configuration for an agent""" name: str model: str temperature: float = 0.2 max_tokens: int = 1000 timeout_ms: int = 30000 max_retries: int = 2 @dataclass class AgentResult: """Result from agent execution""" agent_name: str success: bool data: Dict[str, Any] error: Optional[str] = None tokens_used: Dict[str, int] = None execution_time_ms: float = 0 timestamp: datetime = None def __post_init__(self): if self.timestamp is None: self.timestamp = datetime.now() class Agent(ABC): """ Base class for all copilot agents """ def __init__(self, config: AgentConfig, api_pool): """ Initialize agent Args: config: AgentConfig with model, temperature, etc api_pool: GroqAPIPool for managing API keys """ self.config = config self.api_pool = api_pool self.logger = logger @abstractmethod async def invoke(self, **inputs) -> AgentResult: """ Execute agent with given inputs Must be implemented by subclass Returns: AgentResult with data, success, error """ pass @abstractmethod def _build_prompt(self, **inputs) -> str: """ Build prompt for the agent Must be implemented by subclass """ pass @abstractmethod def _parse_response(self, response_text: str) -> Dict[str, Any]: """ Parse LLM response into structured format Must be implemented by subclass """ pass async def _call_groq(self, prompt: str, model: str = None, conversation_id: str = None) -> str: """ Internal: Call Groq API with prompt """ if model is None: model = self.config.model try: # Try with agent_name (for Gemini pool) or fall back without try: response = await self.api_pool.invoke( prompt=prompt, model=model, temperature=self.config.temperature, max_tokens=self.config.max_tokens, agent_name=self.config.name, conversation_id=conversation_id, ) except TypeError: # Older pool interface (Groq pool doesn't take agent_name yet) response = await self.api_pool.invoke( prompt=prompt, model=model, temperature=self.config.temperature, max_tokens=self.config.max_tokens, ) return response except asyncio.TimeoutError: raise Exception(f"Agent {self.config.name} timeout after {self.config.timeout_ms}ms") except Exception as e: raise Exception(f"Groq API error: {str(e)}") async def _call_with_retry(self, func, max_retries: int = None) -> Any: """ Call function with exponential backoff retry """ if max_retries is None: max_retries = self.config.max_retries for attempt in range(max_retries): try: return await func() except Exception as e: if attempt == max_retries - 1: raise wait_time = 2 ** attempt self.logger.warning( f"Agent {self.config.name} attempt {attempt + 1} failed, " f"retrying in {wait_time}s: {str(e)}" ) await asyncio.sleep(wait_time) def _safe_json_parse(self, text: str, fallback: Dict = None) -> Dict: """ Safely parse JSON from LLM response Falls back to empty dict if parsing fails """ try: # Try to extract JSON from response # Sometimes LLM adds markdown code blocks if "```json" in text: json_str = text.split("```json")[1].split("```")[0] elif "```" in text: json_str = text.split("```")[1].split("```")[0] else: json_str = text return json.loads(json_str) except (json.JSONDecodeError, IndexError, AttributeError): self.logger.warning("Failed to parse JSON response, using fallback") return fallback or {} def _estimate_tokens(self, text: str) -> int: """ Rough estimate of tokens in text 1 token ≈ 4 characters (rough heuristic) """ return max(1, len(text) // 4) async def _create_result( self, success: bool, data: Dict[str, Any], error: str = None, tokens_input: int = 0, tokens_output: int = 0, start_time: float = None, ) -> AgentResult: """ Create AgentResult with proper metadata """ execution_time = 0 if start_time: execution_time = (time.time() - start_time) * 1000 # Convert to ms return AgentResult( agent_name=self.config.name, success=success, data=data, error=error, tokens_used={ "input": tokens_input, "output": tokens_output, "total": tokens_input + tokens_output, }, execution_time_ms=execution_time, ) class SimpleAgent(Agent): """ Simple agent that calls Groq and parses JSON response Subclass this for basic agents that don't need custom logic """ async def invoke(self, **inputs) -> AgentResult: """ Standard invoke flow for simple agents """ start_time = time.time() try: # Build prompt prompt = self._build_prompt(**inputs) tokens_input = self._estimate_tokens(prompt) # Call Groq response = await self._call_groq(prompt) tokens_output = self._estimate_tokens(response) # Parse response parsed_data = self._parse_response(response) # Return success result = await self._create_result( success=True, data=parsed_data, tokens_input=tokens_input, tokens_output=tokens_output, start_time=start_time, ) self.logger.info( f"[OK] {self.config.name}: success in {result.execution_time_ms:.0f}ms " f"({tokens_input + tokens_output} tokens)" ) return result except Exception as e: error_msg = str(e) self.logger.error(f"[FAIL] {self.config.name}: {error_msg}") result = await self._create_result( success=False, data={}, error=error_msg, start_time=start_time, ) return result # Export __all__ = ["Agent", "SimpleAgent", "AgentConfig", "AgentResult"]