From b0422c9e1e72297088f3a3c4379ad6879bf53e65 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?=
 =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= <kirilinka320@gmail.com>
Date: Thu, 4 Jun 2026 23:23:32 +0000
Subject: [PATCH] Add src/agent.py

---
 src/agent.py | 84 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 84 insertions(+)
 create mode 100644 src/agent.py

diff --git a/src/agent.py b/src/agent.py
new file mode 100644
index 0000000..108f955
--- /dev/null
+++ b/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()
\ No newline at end of file