Agent Router in Machine Learning

from semantic_router import Route, RouteLayer
from semantic_router.encoders import OpenAIEncoder

# Define routes with example utterances for each agent
sql_agent_route = Route(
    name="sql_agent",
    utterances=[
        "Show me the top 10 customers by revenue",
        "How many orders were placed last month?",
        "Query the database for active subscriptions",
        "Get the average order value by city",
        "Write a SQL query to find duplicate records",
    ],
)

search_agent_route = Route(
    name="search_agent",
    utterances=[
        "Find documents about machine learning best practices",
        "Search for information about RAG pipelines",
        "What does our documentation say about deployment?",
        "Look up the API reference for authentication",
        "Find relevant papers on transformer architectures",
    ],
)

code_agent_route = Route(
    name="code_agent",
    utterances=[
        "Write a Python function to parse JSON",
        "Debug this error in my code",
        "Refactor this function to be more efficient",
        "Generate unit tests for the payment module",
        "Review this pull request for security issues",
    ],
)

chitchat_route = Route(
    name="chitchat_agent",
    utterances=[
        "Hello, how are you?",
        "What can you help me with?",
        "Thanks for your help",
        "Good morning",
        "Tell me a joke",
    ],
)

# Initialize the encoder and route layer
encoder = OpenAIEncoder(name="text-embedding-3-small")
route_layer = RouteLayer(
    encoder=encoder,
    routes=[sql_agent_route, search_agent_route, code_agent_route, chitchat_route],
)

# Route a query
result = route_layer("How many users signed up this week?")
print(f"Routed to: {result.name}")  # Output: sql_agent
print(f"Confidence: {result.similarity_score:.3f}")

# Route an ambiguous query
result = route_layer("Can you help me optimize this database query?")
print(f"Routed to: {result.name}")  # Could be sql_agent or code_agent

from typing import Literal, TypedDict
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
import json

class AgentState(TypedDict):
    query: str
    routed_agent: str
    response: str
    confidence: float

AGENT_REGISTRY = {
    "sql_agent": "Database queries, data analysis, SQL generation.",
    "search_agent": "Documentation search, knowledge base lookup.",
    "code_agent": "Code writing, debugging, refactoring, reviews.",
    "general_agent": "General conversation and fallback.",
}

router_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

def route_query(state: AgentState) -> AgentState:
    """LLM-based routing node."""
    descs = "\n".join(f"- {k}: {v}" for k, v in AGENT_REGISTRY.items())
    response = router_llm.invoke([
        SystemMessage(content=f"Select the best agent.\nAgents:\n{descs}\nRespond JSON: {{\"agent\": \"name\", \"confidence\": 0.0-1.0}}"),
        HumanMessage(content=state["query"]),
    ])
    result = json.loads(response.content)
    return {**state, "routed_agent": result["agent"], "confidence": result["confidence"]}

def confidence_gate(state: AgentState) -> Literal["specialized", "fallback"]:
    return "fallback" if state["confidence"] < 0.7 else "specialized"

def dispatch(state: AgentState) -> str:
    return state["routed_agent"]

# Build the graph
graph = StateGraph(AgentState)
graph.add_node("router", route_query)
for name in AGENT_REGISTRY:
    graph.add_node(name, lambda s, n=name: {**s, "response": f"[{n}] {s['query']}"})

graph.set_entry_point("router")
graph.add_conditional_edges("router", confidence_gate, {"specialized": "dispatch", "fallback": "general_agent"})
graph.add_node("dispatch", lambda s: s)
graph.add_conditional_edges("dispatch", dispatch, {n: n for n in AGENT_REGISTRY})
for name in AGENT_REGISTRY:
    graph.add_edge(name, END)

app = graph.compile()
result = app.invoke({"query": "How many active users?", "routed_agent": "", "response": "", "confidence": 0.0})
print(f"Routed to: {result['routed_agent']} (confidence: {result['confidence']:.2f})")

import re
from dataclasses import dataclass, field
from typing import Optional
import numpy as np
from sentence_transformers import SentenceTransformer

@dataclass
class RoutingResult:
    agent_id: str
    confidence: float
    method: str  # 'rule', 'semantic', 'llm', 'fallback'
    latency_ms: float = 0.0

@dataclass
class AgentRoute:
    agent_id: str
    description: str
    utterances: list[str] = field(default_factory=list)
    keywords: list[str] = field(default_factory=list)
    patterns: list[str] = field(default_factory=list)
    is_active: bool = True

class CascadingRouter:
    """Production-grade cascading router: rules -> semantic -> LLM -> fallback."""

    def __init__(
        self,
        routes: list[AgentRoute],
        embedding_model: str = "all-MiniLM-L6-v2",
        semantic_threshold: float = 0.75,
        fallback_agent: str = "general_agent",
    ):
        self.routes = {r.agent_id: r for r in routes}
        self.semantic_threshold = semantic_threshold
        self.fallback_agent = fallback_agent

        # Pre-compute embeddings for semantic routing
        self.encoder = SentenceTransformer(embedding_model)
        self._build_utterance_index()

    def _build_utterance_index(self):
        """Pre-compute embeddings for all route utterances."""
        self.utterance_embeddings = {}
        for agent_id, route in self.routes.items():
            if route.utterances:
                self.utterance_embeddings[agent_id] = self.encoder.encode(
                    route.utterances, normalize_embeddings=True
                )

    def route(self, query: str) -> RoutingResult:
        """Cascade through routing strategies."""
        import time

        start = time.perf_counter()

        # Layer 1: Rule-based routing (fastest)
        result = self._rule_based_route(query)
        if result:
            result.latency_ms = (time.perf_counter() - start) * 1000
            return result

        # Layer 2: Semantic routing
        result = self._semantic_route(query)
        if result and result.confidence >= self.semantic_threshold:
            result.latency_ms = (time.perf_counter() - start) * 1000
            return result

        # Layer 3: LLM-based routing (most expensive)
        result = self._llm_route(query)
        if result and result.confidence >= 0.6:
            result.latency_ms = (time.perf_counter() - start) * 1000
            return result

        # Layer 4: Fallback
        elapsed = (time.perf_counter() - start) * 1000
        return RoutingResult(
            agent_id=self.fallback_agent,
            confidence=0.0,
            method="fallback",
            latency_ms=elapsed,
        )

    def _rule_based_route(self, query: str) -> Optional[RoutingResult]:
        """Check keyword and regex patterns."""
        query_lower = query.lower()
        for agent_id, route in self.routes.items():
            # Keyword matching
            for kw in route.keywords:
                if kw.lower() in query_lower:
                    return RoutingResult(agent_id=agent_id, confidence=1.0, method="rule")
            # Regex patterns
            for pattern in route.patterns:
                if re.search(pattern, query, re.IGNORECASE):
                    return RoutingResult(agent_id=agent_id, confidence=1.0, method="rule")
        return None

    def _semantic_route(self, query: str) -> Optional[RoutingResult]:
        """Embedding similarity routing."""
        query_emb = self.encoder.encode([query], normalize_embeddings=True)[0]
        best_agent, best_score = None, -1.0

        for agent_id, embs in self.utterance_embeddings.items():
            scores = embs @ query_emb  # cosine similarity (normalized)
            max_score = float(np.max(scores))
            if max_score > best_score:
                best_score = max_score
                best_agent = agent_id

        if best_agent:
            return RoutingResult(agent_id=best_agent, confidence=best_score, method="semantic")
        return None

    def _llm_route(self, query: str) -> Optional[RoutingResult]:
        """LLM-based routing for ambiguous queries."""
        from openai import OpenAI
        import json

        client = OpenAI()
        descriptions = "\n".join(
            f"- {r.agent_id}: {r.description}"
            for r in self.routes.values()
            if r.is_active
        )

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            temperature=0,
            response_format={"type": "json_object"},
            messages=[
                {"role": "system", "content": f"Select the best agent for the query. Agents:\n{descriptions}\nRespond as JSON: {{\"agent\": \"id\", \"confidence\": 0.0-1.0}}"},
                {"role": "user", "content": query},
            ],
        )
        result = json.loads(response.choices[0].message.content)
        return RoutingResult(
            agent_id=result["agent"],
            confidence=result["confidence"],
            method="llm",
        )


# Usage
routes = [
    AgentRoute(
        agent_id="sql_agent",
        description="Handles database queries and data analysis",
        utterances=["Show me revenue data", "Count active users", "Query the orders table"],
        keywords=["sql", "query", "database"],
        patterns=[r"how many .*(users|orders|customers)"],
    ),
    AgentRoute(
        agent_id="refund_agent",
        description="Processes refund requests and payment issues",
        utterances=["I want a refund", "My payment failed", "Cancel my subscription"],
        keywords=["refund", "cancel", "payment failed"],
        patterns=[r"(refund|cancel).*(order|subscription|payment)"],
    ),
    AgentRoute(
        agent_id="search_agent",
        description="Searches knowledge base and documentation",
        utterances=["Find docs about APIs", "Search for deployment guide", "Look up auth docs"],
        keywords=[],
        patterns=[],
    ),
]

router = CascadingRouter(routes=routes, semantic_threshold=0.72)
result = router.route("I paid for my order but want my money back")
print(f"Agent: {result.agent_id}, Method: {result.method}, Confidence: {result.confidence:.2f}")
# Output: Agent: refund_agent, Method: semantic, Confidence: 0.85

import json
from dataclasses import dataclass
from collections import defaultdict

@dataclass
class RoutingTestCase:
    query: str
    expected_agent: str
    category: str = "general"  # for per-category analysis

def evaluate_router(router, test_cases: list[RoutingTestCase]) -> dict:
    """Evaluate routing accuracy with detailed breakdown."""
    results = {
        "total": len(test_cases),
        "correct": 0,
        "incorrect": 0,
        "fallback_count": 0,
        "method_distribution": defaultdict(int),
        "per_agent_accuracy": defaultdict(lambda: {"correct": 0, "total": 0}),
        "per_category_accuracy": defaultdict(lambda: {"correct": 0, "total": 0}),
        "latencies_ms": [],
        "misroutes": [],
    }

    for tc in test_cases:
        result = router.route(tc.query)
        is_correct = result.agent_id == tc.expected_agent

        if is_correct:
            results["correct"] += 1
        else:
            results["incorrect"] += 1
            results["misroutes"].append({
                "query": tc.query,
                "expected": tc.expected_agent,
                "actual": result.agent_id,
                "confidence": result.confidence,
                "method": result.method,
            })

        if result.method == "fallback":
            results["fallback_count"] += 1

        results["method_distribution"][result.method] += 1
        results["per_agent_accuracy"][tc.expected_agent]["total"] += 1
        results["per_category_accuracy"][tc.category]["total"] += 1
        if is_correct:
            results["per_agent_accuracy"][tc.expected_agent]["correct"] += 1
            results["per_category_accuracy"][tc.category]["correct"] += 1
        results["latencies_ms"].append(result.latency_ms)

    # Compute summary metrics
    results["accuracy"] = results["correct"] / results["total"]
    results["fallback_rate"] = results["fallback_count"] / results["total"]
    results["avg_latency_ms"] = sum(results["latencies_ms"]) / len(results["latencies_ms"])
    results["p95_latency_ms"] = sorted(results["latencies_ms"])[int(0.95 * len(results["latencies_ms"]))]

    for agent, stats in results["per_agent_accuracy"].items():
        stats["accuracy"] = stats["correct"] / stats["total"] if stats["total"] > 0 else 0

    return results


# Example test suite
test_cases = [
    RoutingTestCase("How many orders were placed today?", "sql_agent", "data"),
    RoutingTestCase("Show me the revenue breakdown by state", "sql_agent", "data"),
    RoutingTestCase("I want my money back for order #12345", "refund_agent", "support"),
    RoutingTestCase("Cancel my subscription immediately", "refund_agent", "support"),
    RoutingTestCase("Find the API documentation for payments", "search_agent", "search"),
    RoutingTestCase("What does our wiki say about deployments?", "search_agent", "search"),
    # Edge cases - ambiguous queries
    RoutingTestCase("The payment query is failing", "sql_agent", "ambiguous"),
    RoutingTestCase("Help me with the refund process documentation", "search_agent", "ambiguous"),
]

# Run evaluation
# results = evaluate_router(router, test_cases)
# print(json.dumps(results, indent=2, default=str))

# Agent Router Configuration (YAML)
router:
  strategy: cascading  # rule -> semantic -> llm -> fallback
  
  rule_engine:
    enabled: true
    priority: 1
  
  semantic_engine:
    enabled: true
    priority: 2
    model: all-MiniLM-L6-v2
    threshold: 0.75
    max_latency_ms: 50
  
  llm_engine:
    enabled: true
    priority: 3
    model: gpt-4o-mini
    temperature: 0
    threshold: 0.6
    max_latency_ms: 1000
    max_retries: 2
  
  fallback:
    agent_id: general_agent
    log_level: warn
  
  load_balancer:
    strategy: least-connections  # round-robin | least-connections | latency-based
    health_check_interval_s: 30
    circuit_breaker:
      failure_threshold: 5
      recovery_timeout_s: 60
  
agent_registry:
  - id: sql_agent
    description: "Database queries and data analysis"
    utterances_file: routes/sql_agent.txt
    keywords: [sql, query, database, table, count]
    max_concurrency: 10
    timeout_s: 30
    
  - id: search_agent
    description: "Knowledge base and document search"
    utterances_file: routes/search_agent.txt
    keywords: [search, find, lookup, documentation]
    max_concurrency: 20
    timeout_s: 15
    
  - id: refund_agent
    description: "Refund processing and payment issues"
    utterances_file: routes/refund_agent.txt
    keywords: [refund, cancel, payment, money back]
    max_concurrency: 5
    timeout_s: 45

monitoring:
  log_all_routing_decisions: true
  alert_on_fallback_rate_above: 0.10
  alert_on_accuracy_below: 0.93
  eval_dataset: tests/routing_eval.json
  eval_schedule: daily