src/llm/agent_capabilities.py

from dataclasses import dataclass
from typing import Dict, Literal, Optional


@dataclass
class AgentCapabilityRequirements:
    tools_required: bool
    strict_json_required: bool
    latency_preference: Optional[Literal["fast", "medium", "slow"]] = None


AGENT_CAPABILITIES: Dict[str, AgentCapabilityRequirements] = {
    "InsightsAgent": AgentCapabilityRequirements(tools_required=False, strict_json_required=True),
    "PlanAgent": AgentCapabilityRequirements(tools_required=False, strict_json_required=True),
    "VisualizationAgent": AgentCapabilityRequirements(
        tools_required=True, strict_json_required=True
    ),
    "Router": AgentCapabilityRequirements(
        tools_required=False, strict_json_required=True, latency_preference="fast"
    ),
    "ChatAgent": AgentCapabilityRequirements(tools_required=False, strict_json_required=False),
    "BriefService": AgentCapabilityRequirements(
        tools_required=False, strict_json_required=False, latency_preference="fast"
    ),
}