Add src/agent.py

src/agent.py

@@ -0,0 +1,84 @@
+"""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()