add main.py

main.py

@@ -1,58 +1,38 @@
 """FAQ Bot with ChromaDB and a mock MCP tool.
 
 The agent answers questions about course materials using a local Chroma vector store.
-If the question is about course metadata (e.g., schedule), it calls a simple HTTP
-mock that returns JSON.  The agent is built with LangGraph.
+If the question is about course metadata (e.g., schedule), it calls a simple HTTP mock that returns JSON.  The agent is built with LangGraph.
 
 Run with:
     python main.py
 
-The script will load the FAQ files, build the vector store, and then
-enter an interactive loop.
+The script will load the FAQ files, build the vector store, and then enter an interactive loop.
 """
 
-import os
 import json
-import httpx
+import pathlib
 from pathlib import Path
-from dotenv import load_dotenv
 
 from langchain_ollama import ChatOllama
 from langchain_chroma import Chroma
-from langchain_core.prompts import ChatPromptTemplate
+from langchain_text_splitters import MarkdownHeaderTextSplitter
+from langchain_community.document_loaders import TextLoader
+from langchain_community.embeddings import OllamaEmbeddings
 from langgraph.prebuilt import create_react_agent
 from langgraph.graph import StateGraph, MessagesState
 
-# Load environment variables (e.g., Ollama host)
-load_dotenv()
-
 # ---------- 1. Load FAQ files and create Chroma store ----------
 
-def load_faq_to_chroma(data_dir: str = "data", persist_dir: str = "./chroma_faq"):
-    """Read all .md files in data_dir, chunk them, and persist to Chroma.
-    Returns the Chroma vector store.
-    """
-    from langchain_text_splitters import MarkdownHeaderTextSplitter
-    from langchain_community.document_loaders import TextLoader
-    from langchain_community.embeddings import OllamaEmbeddings
-
-    # Ensure persist directory exists
+def load_faq_to_chroma(data_dir: str = "data", persist_dir: str = "./chroma_faq") -> Chroma:
+    """Read all .md files in data_dir, chunk them, and persist to Chroma."""
     Path(persist_dir).mkdir(parents=True, exist_ok=True)
-
-    # Load all markdown files
     docs = []
     for md_file in Path(data_dir).glob("*.md"):
         loader = TextLoader(str(md_file), encoding="utf-8")
         docs.extend(loader.load())
-
-    # Split documents by headers for better context
     splitter = MarkdownHeaderTextSplitter(headers_to_split_on=["#", "##", "###"])
     split_docs = splitter.split_documents(docs)
-
-    # Use Ollama embeddings
     embeddings = OllamaEmbeddings(model="nomic-embed-text")
-
-    # Persist to Chroma
     chroma = Chroma.from_documents(
         documents=split_docs,
         embedding=embeddings,
@@ -67,17 +47,9 @@ def fetch_course_meta(query: str) -> dict:
     In production this would be an HTTP call to an MCP server.
     Here we simulate with a local JSON file.
     """
-    # For demo purposes, we load a static JSON file.
-    meta_path = Path("course_meta.json")
+    meta_path = Path("mock_meta.json")
     if not meta_path.exists():
-        # Create a simple default meta file if missing
-        meta = {
-            "schedule": {
-                "Monday": "10:00-12:00",
-                "Wednesday": "14:00-16:00",
-            },
-            "instructor": "Prof. Smith",
-        }
+        meta = {"schedule": {"Monday": "10:00-12:00", "Wednesday": "14:00-16:00"}, "instructor": "Prof. Smith"}
         meta_path.write_text(json.dumps(meta, indent=2))
     else:
         meta = json.loads(meta_path.read_text())
@@ -88,22 +60,10 @@ def fetch_course_meta(query: str) -> dict:
 
 def build_agent(chroma: Chroma):
     """Create a LangGraph agent that routes to either Chroma or the MCP tool."""
-    # Define the tool for metadata
     def meta_tool(query: str):
         return json.dumps(fetch_course_meta(query), indent=2)
-
-    # Register tool
     tools = {"fetch_course_meta": meta_tool}
-
-    # Prompt template with source hint
-    prompt = ChatPromptTemplate.from_messages([
-        ("system", "You are an FAQ bot. Use the provided tools wisely.")
-    ])
-
-    # Create a simple React agent with tool calling
-    agent = create_react_agent(ChatOllama(model="llama3"), tools)
-
-    # Graph state: messages only
+    agent = create_react_agent(ChatOllama(model="nomic-embed-text"), tools)
     graph = StateGraph(MessagesState)
     graph.add_node("agent", agent)
     graph.set_entry_point("agent")
@@ -114,26 +74,24 @@ def build_agent(chroma: Chroma):
 def main():
     chroma = load_faq_to_chroma()
     agent = build_agent(chroma)
-
-    print("FAQ Bot ready. Type your question (or 'exit' to quit).")
+    sample_questions = [
+        "What is the deadline for the assignment?",
+        "How many chapters are in the course?",
+        "What is the schedule for the next lecture?",
+    ]
+    for q in sample_questions:
+        print("\nQuestion:", q)
+        response = agent.invoke({"messages": [{"role": "user", "content": q}]})
+        print("Answer:", response["messages"][0]["content"])
     while True:
-        user_input = input("> ")
-        if user_input.lower() in {"exit", "quit"}:
+        try:
+            user_q = input("\nAsk a question (or 'exit'): ")
+        except EOFError:
             break
-        # Determine if question is about metadata by simple keyword check
-        if any(k in user_input.lower() for k in ["schedule", "instructor", "date"]):
-            # Route to MCP tool
-            response = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
-        else:
-            # Route to Chroma via retrieval
-            # Retrieve top k docs
-            docs = chroma.similarity_search(user_input, k=3)
-            context = "\n---\n".join(doc.page_content for doc in docs)
-            # Ask the model with context
-            llm = ChatOllama(model="llama3")
-            answer = llm.invoke([{"role": "system", "content": "You are an FAQ bot."},
-                                 {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {user_input}"}])
-            print(answer.content)
+        if user_q.lower() in {"exit", "quit"}:
+            break
+        response = agent.invoke({"messages": [{"role": "user", "content": user_q}]})
+        print("Answer:", response["messages"][0]["content"])
 
 if __name__ == "__main__":
     main()