84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
"""LangGraph agent that routes queries to either Chroma or the MCP tool.
|
||
|
||
The agent uses a simple decision rule: if the query contains the word
|
||
"schedule" or "metadata", it calls the MCP tool; otherwise it searches
|
||
Chroma. The answer is annotated with a `source` field.
|
||
"""
|
||
|
||
from typing import Dict
|
||
|
||
from langgraph.graph import StateGraph, END
|
||
from langchain_core.runnables import Runnable
|
||
|
||
from .rag_tools import search_course_docs
|
||
from .tools import fetch_course_meta
|
||
|
||
# ---------- State definition ----------
|
||
class AgentState(dict):
|
||
"""State passed between nodes.
|
||
|
||
Keys:
|
||
- ``input``: the user question
|
||
- ``answer``: final answer string
|
||
- ``source``: "chroma" or "mcp_meta"
|
||
"""
|
||
|
||
def __init__(self, input: str = ""):
|
||
super().__init__(input=input, answer="", source="")
|
||
|
||
# ---------- Decision node ----------
|
||
async def decide_route(state: AgentState) -> AgentState:
|
||
"""Return the state unchanged – routing is handled by a conditional edge.
|
||
|
||
The decision logic is intentionally simple to keep the example
|
||
readable.
|
||
"""
|
||
return state
|
||
|
||
# ---------- Chroma search node ----------
|
||
async def search_chroma(state: AgentState) -> AgentState:
|
||
query = state["input"]
|
||
docs = search_course_docs(query)
|
||
# Concatenate the snippets – in a real system we might use a
|
||
# retrieval‑augmented generation chain.
|
||
snippet = "\n\n".join(doc.page_content for doc in docs)
|
||
state["answer"] = snippet if snippet else "No relevant information found."
|
||
state["source"] = "chroma"
|
||
return state
|
||
|
||
# ---------- MCP tool node ----------
|
||
async def fetch_meta(state: AgentState) -> AgentState:
|
||
query = state["input"]
|
||
meta = fetch_course_meta(query)
|
||
state["answer"] = meta if isinstance(meta, str) else str(meta)
|
||
state["source"] = "mcp_meta"
|
||
return state
|
||
|
||
# ---------- Final node ----------
|
||
async def finalize(state: AgentState) -> AgentState:
|
||
# The agent already has the answer; we just return the state.
|
||
return state
|
||
|
||
# ---------- Build the graph ----------
|
||
async def create_agent_executor() -> Runnable:
|
||
# Build the graph
|
||
graph = StateGraph(AgentState)
|
||
graph.add_node("decide", decide_route)
|
||
graph.add_node("search_chroma", search_chroma)
|
||
graph.add_node("fetch_meta", fetch_meta)
|
||
graph.add_node("finalize", finalize)
|
||
|
||
graph.set_entry_point("decide")
|
||
# Conditional edge based on the query content
|
||
def decide_condition(state: AgentState) -> str:
|
||
query = state["input"].lower()
|
||
if "schedule" in query or "metadata" in query:
|
||
return "fetch_meta"
|
||
return "search_chroma"
|
||
|
||
graph.add_conditional_edges("decide", decide_condition)
|
||
graph.add_edge("search_chroma", "finalize")
|
||
graph.add_edge("fetch_meta", "finalize")
|
||
graph.add_edge("finalize", END)
|
||
|
||
return graph.compile() |