Source code for core_genai.interfaces.agent

# -*- coding: utf-8 -*-

"""
Base interfaces (sync & async) for all GenAI agent
implementations.
"""

from abc import abstractmethod
from typing import Any, List, TypedDict, Dict, ClassVar

from core_mixins.interfaces.factory import IFactory




class UsageMetadata(TypedDict, total=False):
    """Normalized token usage returned by every agent's get_metadata."""

    input_tokens: int
    output_tokens: int
    total_tokens: int

    cached_tokens: int
    """Cache read/hit tokens."""

    cache_creation_tokens: int
    """Cache write tokens. Claude only."""

    reasoning_tokens: int
    """Thinking/reasoning tokens. Gemini and Grok only."""

    cost_usd: float
    """Estimated total cost in USD."""





class IAgent(IFactory):
    """Interface class for all agents."""

    # Each concrete must define its own prices...
    _PRICING: ClassVar[Dict[str, Dict[str, float]]] = {}

    _api_key: str
    _client: Any
    _kwargs: Any



    def __init__(self, api_key: str, **kwargs) -> None:
        self._client = None
        self._api_key = api_key
        self._kwargs = kwargs




    @classmethod
    def registration_key(cls) -> str:
        return cls.__name__




    @classmethod
    def create_agent(cls, cls_name: str, api_key: str, **kwargs) -> Any:
        """Instantiates and returns an agent by its registered class name."""

        _cls = cls.get_class(cls_name)
        if _cls is None:
            raise ValueError(f"Agent type {cls_name} not found")
        return _cls(api_key, **kwargs)


    @property
    def client(self) -> Any:
        """Must return the client object."""

        if self._client is None:
            self._client = self._create_client()
        return self._client

    @abstractmethod
    def _create_client(self) -> Any:
        """Must implement and return the specific client."""



    @abstractmethod
    async def analyze(self, model: str, prompt: Any, **kwargs) -> Any:
        """Given a model and prompt the concrete agent must do the analysis."""




    @abstractmethod
    def get_text(self, output: Any) -> str:
        """Given a model output, it retrieves the response (text information)."""




    @abstractmethod
    def get_texts(self, output: Any) -> List[str]:
        """Given a model output, it retrieves all completion choices/candidates."""


    def _get_pricing(self, model: str) -> Dict[str, float]:
        if model in self._PRICING:
            return self._PRICING[model]

        for key, pricing in self._PRICING.items():
            if model.startswith(key):
                return pricing

        return {"input": 0.0, "output": 0.0}



    @abstractmethod
    def get_cost(self, output: Any) -> float:
        """Returns the estimated cost in USD for the given response."""




    @abstractmethod
    def get_metadata(self, output: Any) -> UsageMetadata:
        """
        Given a model output, it retrieves the
        metadata (metadata, cost, etc.) of request.
        """