Solution published: update src/agent.py

src/agent.py

@@ -1,57 +1,64 @@
-import argparse
-from langchain.agents.openai_functions import create_openai_functions_agent
-from langchain.agents import AgentExecutor
-from langchain_core.prompts import ChatPromptTemplate
+import os
+from typing import Dict
+
+from langchain.agents import AgentExecutor, create_openai_tools_agent
 from langchain_ollama import Ollama
-from src.utils import load_faq_to_chroma, search_course_docs, fetch_course_meta
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.messages import HumanMessage
 
-# Initialize embeddings and LLM
-llm = Ollama(model="llama3.1")
-
-# Load or create Chroma collection
-try:
-    chroma = load_faq_to_chroma()
-except Exception:
-    chroma = None
+from src.utils import search_course_docs
+from src.mcp_tool import fetch_course_meta
 
 # Define tools
-from langchain.tools import tool
 
-@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])
+def chroma_tool(query: str):
+    return search_course_docs(query)
 
-@tool
-def fetch_course_meta_tool(query: str) -> str:
-    """Fetch course metadata via MCP-style tool."""
-    results = fetch_course_meta(query)
-    return str(results)
+# Wrap MCP tool
 
-tools = [search_course_docs_tool, fetch_course_meta_tool]
+def mcp_tool(query: str):
+    return fetch_course_meta(query)
 
-# 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.")
+# Simple tool registry
+TOOL_DEFINITIONS = {
+    "search_course_docs": {
+        "func": chroma_tool,
+        "description": "Search FAQ documents in Chroma. Use when question about course content.",
+    },
+    "fetch_course_meta": {
+        "func": mcp_tool,
+        "description": "Get course metadata such as schedule. Use when question about schedule or instructor.",
+    },
+}
+
+# Prompt template
+SYSTEM_PROMPT = """You are a FAQ bot for the course. Use only the tools provided. Do not call both tools unless necessary. Indicate source in your answer: 'source: chroma' or 'source: mcp_meta'."""
+
+PROMPT = ChatPromptTemplate.from_messages([
+    ("system", SYSTEM_PROMPT),
+    ("human", "{input}"),
 ])
 
-agent = create_openai_functions_agent(llm=llm, tools=tools, prompt=prompt)
-executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+# Create agent
+llm = Ollama(model="llama2")
+agent = create_openai_tools_agent(llm, TOOL_DEFINITIONS, PROMPT)
+executor = AgentExecutor(agent=agent, tools=list(TOOL_DEFINITIONS.values()), verbose=True)
 
+# CLI entry point
 if __name__ == "__main__":
+    import argparse
     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"])
+        result = executor.invoke({"input": args.question})
+        print(result["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"])
+        print("Enter questions (Ctrl-D to exit):")
+        try:
+            while True:
+                q = input("Q: ")
+                res = executor.invoke({"input": q})
+                print("A:", res["output"])
+        except EOFError:
+            pass