src/copilot/agents/base.py

"""
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"]