webber/webber-api/src/domains/agents/base.py

"""
Base classes and registry for agent implementations.

All agents are built on PydanticAI and registered in a central registry.
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Protocol, runtime_checkable

from pydantic_ai import Agent

from src.shared.logging import get_logger

logger = get_logger(__name__)


@dataclass
class AgentContext:
    """
    Base context passed to all agent tools.

    Subclass this for agent-specific context (e.g., ExploreContext).
    """
    working_dir: str
    allowed_paths: list[str] = field(default_factory=list)
    timeout_seconds: int = 120


@runtime_checkable
class AgentProtocol(Protocol):
    """Protocol that all agents must implement."""

    @property
    def name(self) -> str:
        """Unique identifier for the agent."""
        ...

    @property
    def description(self) -> str:
        """Human-readable description of what the agent does."""
        ...

    @property
    def agent(self) -> Agent:
        """The underlying PydanticAI agent."""
        ...

    async def run(self, prompt: str, **kwargs: Any) -> str:
        """
        Execute the agent with a prompt.

        Args:
            prompt: User prompt/query
            **kwargs: Additional arguments (working_dir, etc.)

        Returns:
            Agent response as string
        """
        ...


class BaseAgent(ABC):
    """
    Abstract base class for agent implementations.

    Provides common functionality and enforces interface.

    Usage:
        class ExploreAgent(BaseAgent):
            name = "explore"
            description = "Fast codebase exploration"

            def _create_agent(self) -> Agent:
                # Create and configure PydanticAI agent
                ...

            async def run(self, prompt: str, **kwargs) -> str:
                # Execute agent
                ...
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """Unique identifier for the agent."""
        pass

    @property
    @abstractmethod
    def description(self) -> str:
        """Human-readable description."""
        pass

    @property
    def agent(self) -> Agent:
        """Lazy-loaded PydanticAI agent."""
        if not hasattr(self, '_agent') or self._agent is None:
            self._agent = self._create_agent()
        return self._agent

    @abstractmethod
    def _create_agent(self) -> Agent:
        """
        Create and configure the PydanticAI agent.

        Override this to set up model, system prompt, and tools.
        """
        pass

    @abstractmethod
    async def run(self, prompt: str, **kwargs: Any) -> str:
        """Execute the agent."""
        pass


# === Agent Registry ===

_AGENT_REGISTRY: dict[str, BaseAgent] = {}


def register_agent(agent: BaseAgent) -> BaseAgent:
    """
    Register an agent in the global registry.

    Args:
        agent: Agent instance to register

    Returns:
        The registered agent (for decorator chaining)
    """
    if agent.name in _AGENT_REGISTRY:
        logger.warning(f"Overwriting existing agent: {agent.name}")

    _AGENT_REGISTRY[agent.name] = agent
    logger.info(f"Registered agent: {agent.name}")
    return agent


def get_agent(name: str) -> BaseAgent | None:
    """
    Get an agent by name.

    Args:
        name: Agent name

    Returns:
        Agent instance or None if not found
    """
    return _AGENT_REGISTRY.get(name)


def list_agents() -> list[dict[str, str]]:
    """
    List all registered agents.

    Returns:
        List of agent info dicts with name and description
    """
    return [
        {"name": agent.name, "description": agent.description}
        for agent in _AGENT_REGISTRY.values()
    ]


def get_registry() -> dict[str, BaseAgent]:
    """Get the full agent registry."""
    return _AGENT_REGISTRY.copy()