conditional_routing.py

"""Conditional fan-out variant — routes to only the relevant sub-agents.

Uses structured output to classify the query and selectively dispatch
agents instead of always hitting all three.
"""

import operator
from collections.abc import Sequence
from typing import Annotated, TypedDict

from langchain_anthropic import ChatAnthropic
from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.graph import END, START, StateGraph
from langgraph.types import RetryPolicy
from langsmith import traceable, tracing_context
from pydantic import BaseModel, Field

llm = ChatAnthropic(model="claude-sonnet-4-5-20250929", temperature=0)


class State(TypedDict):
    question: str
    research_results: Annotated[list[dict], operator.add]
    final_response: str


# --- Sub-agents (same as parallel_agents.py) ---


def make_agent(name: str, focus: str):
    @traceable(name=name, run_type="chain")
    def node(state: State) -> dict:
        response = llm.invoke([
            SystemMessage(
                content=f"You are the {name} agent. Focus on {focus}. "
                "Return a concise summary. Cite your source type."
            ),
            HumanMessage(content=f"Research query: {state['question']}"),
        ])
        return {"research_results": [{"source": name, "content": response.content}]}

    return node


kb_agent = make_agent("knowledge_base", "internal knowledge base searches.")
web_agent = make_agent("web_search", "recent news and industry trends.")
policy_agent = make_agent("policy", "compliance, legal, and regulatory frameworks.")


@traceable(name="Synthesizer", run_type="chain")
def synthesize(state: State) -> dict:
    context = "\n\n".join(
        f"[{r['source']}]: {r['content']}" for r in state["research_results"]
    )
    response = llm.invoke([
        SystemMessage(
            content="Synthesize the following research into a clear, actionable "
            "response. When policy information conflicts with or constrains "
            "other responses, the policy statement takes precedence. "
            "Never soften or omit policy restrictions."
        ),
        HumanMessage(
            content=f"Customer question: {state['question']}\n\n"
            f"Research findings:\n{context}"
        ),
    ])
    return {"final_response": response.content}


# --- Conditional routing ---


class RoutingPlan(BaseModel):
    agents: list[str] = Field(description="Agents to activate: kb, web, policy")


structured_llm = llm.with_structured_output(RoutingPlan)


@traceable(name="Classifier", run_type="chain")
def classify_and_route(state: State) -> Sequence[str]:
    """Classify the query and return which agents to invoke."""
    plan = structured_llm.invoke([
        SystemMessage(
            content="Decide which research agents to invoke. "
            "Available: kb, web, policy. When in doubt, include the agent."
        ),
        HumanMessage(content=state["question"]),
    ])
    return plan.agents or ["kb"]


# --- Graph with conditional fan-out ---

builder = StateGraph(State)

builder.add_node("kb", kb_agent, retry=RetryPolicy(max_attempts=3))
builder.add_node("web", web_agent, retry=RetryPolicy(max_attempts=3))
builder.add_node("policy", policy_agent, retry=RetryPolicy(max_attempts=3))
builder.add_node("synthesize", synthesize)

# Conditional fan-out from START based on classifier
builder.add_conditional_edges(START, classify_and_route, ["kb", "web", "policy"])

# Individual edges from each agent to synthesize (not list-style fan-in,
# because with conditional routing some branches may not be dispatched)
builder.add_edge("kb", "synthesize")
builder.add_edge("web", "synthesize")
builder.add_edge("policy", "synthesize")
builder.add_edge("synthesize", END)

graph = builder.compile()


if __name__ == "__main__":
    queries = [
        "What is our refund policy for enterprise clients?",
        "How does GDPR affect our data pipeline architecture?",
        "What competitors launched AI features last quarter?",
    ]

    for q in queries:
        print("=" * 60)
        print(f"QUESTION: {q}")
        with tracing_context(
            metadata={"example": True},
            tags=["article-01", "conditional-routing"],
        ):
            result = graph.invoke({"question": q})

        sources = [r["source"] for r in result["research_results"]]
        print(f"Agents dispatched: {sources}")
        print(f"Response: {result['final_response'][:200]}...")
        print()