006b4ee5d7
- Move intent imports to TYPE_CHECKING block in graph.py (TC001) - Rename test classes to CapWords convention (N801) - Fix line length violations across test files (E501) - Auto-fix import sorting (I001)
122 lines
3.9 KiB
Python
122 lines
3.9 KiB
Python
"""LangGraph Supervisor construction -- connects registry, agents, LLM, and persistence."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import TYPE_CHECKING
|
|
|
|
from langgraph.prebuilt import create_react_agent
|
|
from langgraph_supervisor import create_supervisor
|
|
|
|
from app.agents import get_tools_by_names
|
|
|
|
if TYPE_CHECKING:
|
|
from langchain_core.language_models import BaseChatModel
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
|
from langgraph.graph.state import CompiledStateGraph
|
|
|
|
from app.intent import ClassificationResult, IntentClassifier
|
|
from app.registry import AgentRegistry
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SUPERVISOR_PROMPT = (
|
|
"You are a customer support supervisor. "
|
|
"Route customer requests to the appropriate agent based on their description.\n\n"
|
|
"Available agents and their roles:\n"
|
|
"{agent_descriptions}\n\n"
|
|
"Routing rules:\n"
|
|
"- For order status and tracking queries, use the order_lookup agent.\n"
|
|
"- For order modifications like cancellations, use the order_actions agent.\n"
|
|
"- For discounts, promotions, or coupon codes, use the discount agent.\n"
|
|
"- For anything else or when uncertain, use the fallback agent.\n"
|
|
"- If the user's request involves multiple actions, execute them in order.\n"
|
|
"- If a previous intent classification is provided, follow it.\n"
|
|
)
|
|
|
|
|
|
def _format_agent_descriptions(registry: AgentRegistry) -> str:
|
|
"""Build agent description text for the supervisor prompt."""
|
|
lines = []
|
|
for agent in registry.list_agents():
|
|
lines.append(f"- {agent.name}: {agent.description}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def build_agent_nodes(
|
|
registry: AgentRegistry,
|
|
llm: BaseChatModel,
|
|
) -> list:
|
|
"""Create LangGraph react agent nodes from registry configurations."""
|
|
agent_nodes = []
|
|
for agent_config in registry.list_agents():
|
|
tools = get_tools_by_names(agent_config.tools)
|
|
|
|
system_prompt = (
|
|
f"You are the {agent_config.name} agent. "
|
|
f"Personality: {agent_config.personality.tone}. "
|
|
f"{agent_config.description} "
|
|
f"Permission level: {agent_config.permission}."
|
|
)
|
|
|
|
agent_node = create_react_agent(
|
|
model=llm,
|
|
tools=tools,
|
|
name=agent_config.name,
|
|
prompt=system_prompt,
|
|
)
|
|
agent_nodes.append(agent_node)
|
|
|
|
return agent_nodes
|
|
|
|
|
|
def build_graph(
|
|
registry: AgentRegistry,
|
|
llm: BaseChatModel,
|
|
checkpointer: AsyncPostgresSaver,
|
|
intent_classifier: IntentClassifier | None = None,
|
|
) -> CompiledStateGraph:
|
|
"""Build and compile the LangGraph supervisor graph.
|
|
|
|
If an intent_classifier is provided, the supervisor prompt is enhanced
|
|
with agent descriptions for better routing. The classifier is stored
|
|
for use by the routing layer (ws_handler).
|
|
"""
|
|
agent_nodes = build_agent_nodes(registry, llm)
|
|
agent_descriptions = _format_agent_descriptions(registry)
|
|
|
|
prompt = SUPERVISOR_PROMPT.format(agent_descriptions=agent_descriptions)
|
|
|
|
workflow = create_supervisor(
|
|
agent_nodes,
|
|
model=llm,
|
|
prompt=prompt,
|
|
output_mode="full_history",
|
|
)
|
|
|
|
graph = workflow.compile(checkpointer=checkpointer)
|
|
|
|
# Attach classifier and registry to graph for use by ws_handler
|
|
graph.intent_classifier = intent_classifier # type: ignore[attr-defined]
|
|
graph.agent_registry = registry # type: ignore[attr-defined]
|
|
|
|
return graph
|
|
|
|
|
|
async def classify_intent(
|
|
graph: CompiledStateGraph,
|
|
message: str,
|
|
) -> ClassificationResult | None:
|
|
"""Classify user intent using the graph's attached classifier.
|
|
|
|
Returns None if no classifier is configured.
|
|
"""
|
|
classifier = getattr(graph, "intent_classifier", None)
|
|
registry = getattr(graph, "agent_registry", None)
|
|
|
|
if classifier is None or registry is None:
|
|
return None
|
|
|
|
agents = registry.list_agents()
|
|
return await classifier.classify(message, agents)
|