Initial implementation of FAQ bot with ChromaDB and MCP-style tool: update src/agent.py

src/agent.py

@@ -1,37 +1,57 @@
 import os
-from langchain.agents import initialize_agent, Tool
-from langchain.llms import Ollama
-from langchain.chains import RetrievalQA
-from langchain_ollama import OllamaEmbeddings
-from langchain_chroma import Chroma
-from src.utils import search_course_docs
-from src.mcp_tool import fetch_course_meta, start_meta_server
-
-# Start mock server
-start_meta_server()
+from typing import Dict, Any
+from langchain_ollama import OllamaLLM
+from langchain.agents import initialize_agent, Tool, AgentType
+from langchain.memory import ConversationBufferMemory
+from src.utils import search_course_docs, fetch_course_meta
 
 # Define tools
 search_tool = Tool(
     name="search_course_docs",
-    func=lambda q: "\n".join([doc.page_content for doc in search_course_docs(q, k=3)]),
-    description="Search local FAQ documents. Use when question is about course content."
+    func=search_course_docs,
+    description="Search local FAQ docs in ChromaDB. Use when question about course content. Returns list of relevant documents."
 )
-
 meta_tool = Tool(
     name="fetch_course_meta",
-    func=lambda q: str(fetch_course_meta(q)),
-    description="Get course metadata like schedule or instructor. Use when question is about schedule or meta."
+    func=fetch_course_meta,
+    description="Fetch course metadata (schedule, exams) from MCP-style tool. Use when question about schedule or metadata. Returns list of matching items."
 )
 
-# LLM and agent
-llm = Ollama(model="llama3.1")
-agent = initialize_agent([search_tool, meta_tool], llm, agent_type="zero-shot-react-description", verbose=True)
+# System prompt to guide routing
+SYSTEM_PROMPT = (
+    "You are an FAQ bot for the course. Use search_course_docs for content questions and fetch_course_meta for schedule or metadata questions.\n"
+    "When answering, include a field 'source' with value 'chroma' or 'mcp_meta' to indicate which tool was used."
+)
+
+# LLM and agent setup
+llm = OllamaLLM(model="llama3")
+memory = ConversationBufferMemory(memory_key="chat_history")
+agent = initialize_agent(
+    tools=[search_tool, meta_tool],
+    llm=llm,
+    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+    memory=memory,
+    verbose=True,
+    system_prompt=SYSTEM_PROMPT,
+)
+
+def ask(question: str) -> Dict[str, Any]:
+    response = agent.run(question)
+    # Parse response to extract source if present
+    source = "unknown"
+    if "source:" in response.lower():
+        parts = response.lower().split("source:")
+        source = parts[1].strip().split()[0]
+    return {"answer": response, "source": source}
 
 if __name__ == "__main__":
-    print("FAQ-бот готов. Введите вопрос (или 'exit'): ")
-    while True:
-        q = input("> ")
-        if q.lower() in {"exit", "quit"}:
-            break
-        response = agent.run(q)
-        print(response)
+    # Simple CLI with 3 preset questions
+    questions = [
+        "Как подключить ChromaDB?",
+        "Что такое MCP‑tool?",
+        "Когда проходят экзамены?"
+    ]
+    for q in questions:
+        print("Q:", q)
+        print("A:", ask(q)["answer"], "(source:", ask(q)["source"], ")")
+        print()