web/components/agent_provider_matrix.py

"""Agent-to-routing-policy matrix component for multi-provider configuration."""

from typing import Dict, List, Tuple

import gradio as gr

# List of agents in the system
AGENT_NAMES = [
    "indicator_agent",
    "pattern_agent",
    "trend_agent",
    "decision_agent",
    "fundamentals_agent",
    "sentiment_agent",
    "news_agent",
    "technical_analyst",
    "risk_manager",
    "portfolio_manager",
]

# Routing policy options
ROUTING_POLICIES = [
    ("Auto (Default)", "auto"),
    ("🚀 Fastest Response", ":fastest"),
    ("💰 Cheapest Cost", ":cheapest"),
    ("Groq", "groq"),
    ("Together AI", "together"),
    ("Replicate", "replicate"),
    ("Cerebras", "cerebras"),
    ("Fireworks", "fireworks"),
    ("DeepInfra", "deepinfra"),
    ("Llama 3.3 70B", "meta-llama/Llama-3.3-70B-Instruct"),
]

# Model tier options
MODEL_TIERS = [
    ("Fast (Lower Cost)", "fast"),
    ("Capable (Balanced)", "capable"),
    ("Advanced (Premium)", "vision"),
]


def create_agent_provider_matrix() -> Tuple[gr.Dataframe, Dict]:
    """
    Create agent-to-routing-policy matrix UI component.

    Returns:
        Tuple of (dataframe component, agent_dropdowns dict)
    """
    with gr.Accordion("🎯 Per-Agent Routing Configuration", open=False):
        gr.Markdown("""
### Advanced Multi-Provider Configuration

Assign different routing policies to each agent for optimal cost/performance balance.

**Examples:**
- Fast agents (Indicator, Pattern) → `:cheapest` for cost savings
- Decision-critical agents (Decision, Portfolio Manager) → `:fastest` for best quality
- Research agents (News, Sentiment) → `groq` or specific providers
        """)

        # Create a dictionary to store dropdowns for each agent
        agent_components = {}

        with gr.Column():
            # Header row
            with gr.Row():
                gr.Markdown("**Agent**", elem_classes=["matrix-header"])
                gr.Markdown("**Routing Policy**", elem_classes=["matrix-header"])
                gr.Markdown("**Model Tier**", elem_classes=["matrix-header"])

            # Create a row for each agent
            for agent_name in AGENT_NAMES:
                with gr.Row():
                    # Agent name label
                    display_name = agent_name.replace("_", " ").title()
                    gr.Markdown(f"**{display_name}**")

                    # Routing policy dropdown
                    routing_dropdown = gr.Dropdown(
                        choices=ROUTING_POLICIES,
                        value="auto",
                        label="",
                        show_label=False,
                        container=False,
                    )

                    # Model tier dropdown
                    tier_dropdown = gr.Dropdown(
                        choices=MODEL_TIERS,
                        value="capable",
                        label="",
                        show_label=False,
                        container=False,
                    )

                    agent_components[agent_name] = {
                        "routing_policy": routing_dropdown,
                        "model_tier": tier_dropdown,
                    }

    return agent_components


def get_agent_routing_config(agent_components: Dict) -> Dict[str, Dict[str, str]]:
    """
    Extract current routing configuration from agent components.

    Args:
        agent_components: Dictionary of agent UI components

    Returns:
        Dictionary mapping agent names to their routing config
    """
    config = {}
    for agent_name, components in agent_components.items():
        config[agent_name] = {
            "routing_policy": components["routing_policy"].value,
            "model_tier": components["model_tier"].value,
        }
    return config


def apply_routing_preset(preset: str, agent_components: Dict) -> List[gr.update]:
    """
    Apply a routing preset to all agents.

    Args:
        preset: Preset name ("cost_optimized" or "performance_optimized")
        agent_components: Dictionary of agent UI components

    Returns:
        List of Gradio updates for each dropdown
    """
    updates = []

    if preset == "cost_optimized":
        # All agents use :cheapest routing
        for agent_name in AGENT_NAMES:
            updates.append(gr.update(value=":cheapest"))  # routing policy
            updates.append(gr.update(value="fast"))  # model tier

    elif preset == "performance_optimized":
        # All agents use :fastest routing
        for agent_name in AGENT_NAMES:
            updates.append(gr.update(value=":fastest"))  # routing policy
            updates.append(gr.update(value="vision"))  # model tier

    elif preset == "balanced":
        # Fast agents use :cheapest, decision agents use :fastest
        fast_agents = ["indicator_agent", "pattern_agent", "trend_agent"]
        decision_agents = ["decision_agent", "portfolio_manager", "risk_manager"]

        for agent_name in AGENT_NAMES:
            if agent_name in fast_agents:
                updates.append(gr.update(value=":cheapest"))
                updates.append(gr.update(value="fast"))
            elif agent_name in decision_agents:
                updates.append(gr.update(value=":fastest"))
                updates.append(gr.update(value="vision"))
            else:
                updates.append(gr.update(value="auto"))
                updates.append(gr.update(value="capable"))

    else:
        # Reset to defaults
        for agent_name in AGENT_NAMES:
            updates.append(gr.update(value="auto"))
            updates.append(gr.update(value="capable"))

    return updates


def format_routing_config_status(config: Dict[str, Dict[str, str]]) -> str:
    """
    Format routing configuration status message.

    Args:
        config: Agent routing configuration dictionary

    Returns:
        Formatted status string
    """
    if not config:
        return "Using global routing policy for all agents"

    # Count unique routing policies
    policies = {}
    for agent_name, agent_config in config.items():
        policy = agent_config.get("routing_policy", "auto")
        if policy not in policies:
            policies[policy] = []
        policies[policy].append(agent_name.replace("_", " ").title())

    # Format status
    status_lines = ["✅ Per-Agent Routing Active:"]
    for policy, agents in policies.items():
        if policy.startswith(":"):
            policy_display = policy.upper()
        elif "/" in policy:
            policy_display = policy.split("/")[-1]
        else:
            policy_display = policy.title()

        status_lines.append(f"  • {policy_display}: {len(agents)} agents")

    return "\n".join(status_lines)


def export_routing_config(config: Dict[str, Dict[str, str]]) -> str:
    """
    Export routing configuration as JSON string.

    Args:
        config: Agent routing configuration dictionary

    Returns:
        JSON string of configuration
    """
    import json

    return json.dumps(config, indent=2)


def import_routing_config(config_json: str) -> Dict[str, Dict[str, str]]:
    """
    Import routing configuration from JSON string.

    Args:
        config_json: JSON string of configuration

    Returns:
        Parsed configuration dictionary
    """
    import json

    try:
        config = json.loads(config_json)
        # Validate structure
        if not isinstance(config, dict):
            raise ValueError("Configuration must be a dictionary")

        for agent_name, agent_config in config.items():
            if not isinstance(agent_config, dict):
                raise ValueError(f"Agent config for {agent_name} must be a dictionary")
            if "routing_policy" not in agent_config:
                raise ValueError(f"Missing routing_policy for {agent_name}")

        return config
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON: {str(e)}")