tatlock/src/core/household_registry.py

"""
Household registry for managing agent capabilities and toolsets.

Provides centralized registry of household members (agents) with their
capabilities and tools. Supports two-tier abstraction: executive summaries
for coordination and full toolsets for execution.
"""
from typing import Any, Optional

from pydantic import BaseModel, ConfigDict
from pydantic_ai import Agent

from .logging_config import get_logger

logger = get_logger(__name__)


class HouseholdCapability(BaseModel):
    """
    Executive summary of a household member's capabilities.

    This is what the Steward and Butler see for coordination.
    High-level description without implementation details.
    """
    name: str  # Unique identifier: "tatlock_core", "librarian", "developer"
    role: str  # Display name: "Butler's Core Tools", "The Librarian"
    category: str  # "core", "research", "technical", "automation"
    description: str  # One-sentence description of capabilities
    domains: list[str]  # Capability domains: ["computation", "information", "datetime"]
    cost: str  # "low", "medium", "high" - resource cost estimate
    requires_network: bool  # Whether network access is needed


class HouseholdMember(BaseModel):
    """
    Full specification of a household member.

    Contains both the executive summary (for coordination) and
    implementation details (tools/agent).
    """
    model_config = ConfigDict(arbitrary_types_allowed=True)

    capability: HouseholdCapability
    tools: list[Any]  # PydanticAI tool definitions (any type since Tool is a dataclass)
    agent: Optional[Any] = None  # For expert agents (Phase 4)


class HouseholdRegistry:
    """
    Registry of household capabilities and implementations.

    Manages household members and their tools. Provides:
    1. Executive summaries for Steward/Butler coordination
    2. Full toolsets for scoped execution
    3. Agent delegation (Phase 4)
    """

    def __init__(self):
        """Initialize empty registry."""
        self._members: dict[str, HouseholdMember] = {}
        logger.info("household_registry_initialized")

    def register(
        self,
        name: str,
        capability: HouseholdCapability,
        tools: list[Any],
        agent: Optional[Any] = None,
    ) -> None:
        """
        Register a household member.

        Args:
            name: Unique identifier (must match capability.name)
            capability: Executive summary
            tools: PydanticAI tool definitions
            agent: Optional expert agent for delegation

        Raises:
            ValueError: If name doesn't match capability.name

        Example:
            >>> registry.register(
            ...     name="tatlock_core",
            ...     capability=HouseholdCapability(
            ...         name="tatlock_core",
            ...         role="Butler's Core Tools",
            ...         category="core",
            ...         description="Basic computation, time, and information tools",
            ...         domains=["computation", "datetime", "information"],
            ...         cost="low",
            ...         requires_network=True,
            ...     ),
            ...     tools=[calculator_tool, datetime_tool, search_tool],
            ... )
        """
        if name != capability.name:
            raise ValueError(
                f"Name mismatch: '{name}' != '{capability.name}'"
            )

        self._members[name] = HouseholdMember(
            capability=capability,
            tools=tools,
            agent=agent,
        )

        logger.info(
            "household_member_registered",
            name=name,
            role=capability.role,
            domains=capability.domains,
            tool_count=len(tools),
            has_agent=agent is not None,
        )

    def unregister(self, name: str) -> None:
        """
        Unregister a household member.

        Args:
            name: Member name to remove

        Example:
            >>> registry.unregister("tatlock_core")
        """
        if name in self._members:
            member = self._members.pop(name)
            logger.info(
                "household_member_unregistered",
                name=name,
                role=member.capability.role,
            )

    def get_member(self, name: str) -> Optional[HouseholdMember]:
        """
        Get full household member specification.

        Args:
            name: Member name

        Returns:
            HouseholdMember if found, None otherwise
        """
        return self._members.get(name)

    def get_all_capabilities(self) -> list[HouseholdCapability]:
        """
        Get executive summaries of all household members.

        This is what the Steward sees when analyzing requests.
        Returns high-level capabilities without implementation details.

        Returns:
            List of capability summaries

        Example:
            >>> capabilities = registry.get_all_capabilities()
            >>> for cap in capabilities:
            ...     print(f"{cap.role}: {cap.description}")
        """
        return [member.capability for member in self._members.values()]

    def get_scoped_tools(self, names: list[str]) -> list[Any]:
        """
        Get combined tools from specified household members.

        Creates a scoped toolset containing only tools from
        the requested members. Used to give Tatlock only the
        tools recommended by the Steward.

        Args:
            names: List of member names to include

        Returns:
            Combined list of tool definitions

        Example:
            >>> # Steward recommends only tatlock_core
            >>> tools = registry.get_scoped_tools(["tatlock_core"])
            >>> # Tatlock now has only core tools, not all household tools
        """
        tools = []
        for name in names:
            member = self._members.get(name)
            if member:
                tools.extend(member.tools)
            else:
                logger.warning(
                    "household_member_not_found",
                    requested_name=name,
                    available_names=list(self._members.keys()),
                )

        logger.debug(
            "scoped_tools_created",
            requested_members=names,
            total_tools=len(tools),
        )

        return tools

    def list_members(self) -> list[str]:
        """
        List all registered member names.

        Returns:
            List of member names
        """
        return list(self._members.keys())

    def get_members_by_domain(self, domain: str) -> list[HouseholdCapability]:
        """
        Get capabilities that support a specific domain.

        Args:
            domain: Domain to filter by (e.g., "computation", "research")

        Returns:
            List of capabilities supporting the domain

        Example:
            >>> # Find all members that can do research
            >>> research_caps = registry.get_members_by_domain("research")
        """
        return [
            member.capability
            for member in self._members.values()
            if domain in member.capability.domains
        ]

    def get_members_by_category(self, category: str) -> list[HouseholdCapability]:
        """
        Get capabilities by category.

        Args:
            category: Category to filter by (e.g., "core", "research", "technical")

        Returns:
            List of capabilities in the category
        """
        return [
            member.capability
            for member in self._members.values()
            if member.capability.category == category
        ]

    def __len__(self) -> int:
        """Get number of registered members."""
        return len(self._members)

    def __contains__(self, name: str) -> bool:
        """Check if member is registered."""
        return name in self._members


# Global registry instance
household_registry = HouseholdRegistry()


def get_household_registry() -> HouseholdRegistry:
    """
    Get global household registry instance.

    Returns:
        HouseholdRegistry instance
    """
    return household_registry