Source code for agentopera.router.intent_registry

from typing import Dict, Union, Set, Optional, List
from agentopera.ui.vercel_ai_request import UserAgentInfo
from agentopera.utils.logger import logger



class IntentRegistry:
    """
    Internal configuration holder for intent and agent relationships.

    This class is designed to manage:
    - Intent descriptions for classification models
    - Mapping from intent → agent_id
    - Mapping from agent_id → descriptions

    It acts as the source of truth for the classifier and agent registry.
    """

    def __init__(
        self,
        intent_descriptions: Dict[str, str],
        agent_intent_mapping: Dict[str, str],
        agent_descriptions: Dict[str, Union[str, list[str]]],
        schema_overrides: Optional[Dict[str, dict]] = None
    ):
        """
        Initializes the intent registry.

        Args:
            intent_descriptions: A dictionary of intent → description.
            agent_intent_mapping: A dictionary of intent → agent_id.
            agent_descriptions: A dictionary of agent_id → description(s).
        """
        self.intent_descriptions = intent_descriptions
        self.agent_intent_mapping = agent_intent_mapping
        self.agent_descriptions = agent_descriptions
        self.schema_overrides = schema_overrides or {}
    


    @classmethod
    def from_api_data(cls, api_data: List[UserAgentInfo]) -> 'IntentRegistry':
        intent_descriptions = {}
        agent_descriptions = {}
        agent_intent_mapping = {}

        for agent in api_data:
            # Construct intent ID
            intent_key = f"{agent.agent_type}_{agent.agent_name}_intent"

            # Populate registries
            intent_descriptions[intent_key] = agent.agent_intent
            agent_descriptions[agent.agent_name] = agent.agent_description
            agent_intent_mapping[intent_key] = agent.agent_name

        return IntentRegistry(
            intent_descriptions=intent_descriptions,
            agent_intent_mapping=agent_intent_mapping,
            agent_descriptions=agent_descriptions
        )

    


    def merge(self, other: "IntentRegistry") -> "IntentRegistry":
        conflicts = set(self.intent_descriptions).intersection(other.intent_descriptions)

        for key in conflicts:
            logger.info(f"[IntentRegistry.merge] Overriding intent key '{key}' from default config")

        return IntentRegistry(
            intent_descriptions={**self.intent_descriptions, **other.intent_descriptions},
            agent_intent_mapping={**self.agent_intent_mapping, **other.agent_intent_mapping},
            agent_descriptions={**self.agent_descriptions, **other.agent_descriptions},
            schema_overrides={**self.schema_overrides, **other.schema_overrides},
        )




    def get_all_intents(self) -> Set[str]:
        """Returns all available intent keys."""
        return set(self.intent_descriptions.keys())

    


    def get_schema_override(self, intent: str) -> dict:
        """Returns all overriding schemas if any"""
        return self.schema_overrides.get(intent, {})




    def get_description(self, intent: str) -> str:
        """Returns description for a specific intent."""
        return self.intent_descriptions.get(intent, "No description available")




    def get_agent_for_intent(self, intent: str) -> str:
        """Returns agent_id associated with the given intent."""
        return self.agent_intent_mapping.get(intent, "chat")




    def get_all_agents(self) -> Set[str]:
        """Returns all registered agent IDs."""
        return set(self.agent_descriptions.keys())




    def get_agent_description(self, agent_id: str) -> Union[str, list[str]]:
        """Returns description for the specified agent."""
        return self.agent_descriptions.get(agent_id, "No description available")

    


    def get_tool_schemas(self) -> list[dict]:
        """
        Builds OpenAI-compatible tool definitions for all registered intents.

        Returns:
            List of function tool schema dictionaries for OpenAI-compatible models.
        """
        tools = []

        for intent, description in self.intent_descriptions.items():
            schema_override = self.schema_overrides.get(intent, {})

            default_schema = {
                "type": "object",
                "properties": {
                    "confidence_score": {
                        "type": "string",
                        "description": f"Confidence score for '{intent}'. A string presented Integer between 0 and 100.",
                    }
                },
                "required": ["confidence_score"],
            }

            parameters = {
                **default_schema,
                **schema_override,
                "properties": {
                    **default_schema["properties"],
                    **schema_override.get("properties", {}),
                },
                "required": schema_override.get("required", default_schema["required"]),
            }

            tools.append({
                "type": "function",
                "function": {
                    "name": intent,
                    "description": description,
                    "parameters": parameters,
                }
            })

        return tools