Syncing local state to remote: update src/agent.py

src/agent.py

@@ -1,57 +1,57 @@
-import os
-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
+import argparse
+from langchain.agents import create_openai_functions_agent, AgentExecutor
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.messages import HumanMessage
+from langchain_ollama import Ollama
+from src.utils import load_faq_to_chroma, search_course_docs, fetch_course_meta
+
+# Initialize embeddings and LLM
+llm = Ollama(model="llama3.1")
+
+# Load or create Chroma collection
+try:
+    chroma = load_faq_to_chroma()
+except Exception:
+    chroma = None
 
 # Define tools
-search_tool = Tool(
-    name="search_course_docs",
-    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=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."
-)
+from langchain.tools import tool
 
-# 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."
-)
+@tool
+def search_course_docs_tool(query: str, k: int = 3) -> str:
+    """Search local FAQ docs in ChromaDB."""
+    docs = search_course_docs(query, k)
+    return "\n".join([doc.page_content for doc in docs])
 
-# 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,
-)
+@tool
+def fetch_course_meta_tool(query: str) -> str:
+    """Fetch course metadata via MCP-style tool."""
+    results = fetch_course_meta(query)
+    return str(results)
 
-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}
+tools = [search_course_docs_tool, fetch_course_meta_tool]
+
+# Prompt template with source hint
+prompt = ChatPromptTemplate.from_messages([
+    ("system", "You are a helpful FAQ assistant. Use the tools only when necessary. In your answer, include a line like 'source: chroma' or 'source: mcp_meta' to indicate which tool was used.")
+])
+
+agent = create_openai_functions_agent(llm=llm, tools=tools, prompt=prompt)
+executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
 if __name__ == "__main__":
-    # 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()
+    parser = argparse.ArgumentParser(description="FAQ bot CLI")
+    parser.add_argument("--question", type=str, help="Question to ask the bot")
+    args = parser.parse_args()
+    if args.question:
+        response = executor.invoke({"input": args.question})
+        print(response["output"])
+    else:
+        # Interactive mode
+        print("FAQ Bot. Type 'exit' to quit.")
+        while True:
+            q = input("> ")
+            if q.lower() in ("exit", "quit"):
+                break
+            resp = executor.invoke({"input": q})
+            print(resp["output"])