add main.py

main.py

@@ -0,0 +1,139 @@
+"""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.
+
+Run with:
+    python main.py
+
+The script will load the FAQ files, build the vector store, and then
+enter an interactive loop.
+"""
+
+import os
+import json
+import httpx
+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 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
+    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,
+        persist_directory=persist_dir,
+    )
+    return chroma
+
+# ---------- 2. MCP-style tool: fetch course metadata ----------
+
+def fetch_course_meta(query: str) -> dict:
+    """Mock MCP tool that returns course metadata.
+    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")
+    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_path.write_text(json.dumps(meta, indent=2))
+    else:
+        meta = json.loads(meta_path.read_text())
+    # Return the whole meta; the agent can filter as needed
+    return meta
+
+# ---------- 3. Agent setup ----------
+
+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
+    graph = StateGraph(MessagesState)
+    graph.add_node("agent", agent)
+    graph.set_entry_point("agent")
+    return graph.compile()
+
+# ---------- 4. Interactive CLI ----------
+
+def main():
+    chroma = load_faq_to_chroma()
+    agent = build_agent(chroma)
+
+    print("FAQ Bot ready. Type your question (or 'exit' to quit).")
+    while True:
+        user_input = input("> ")
+        if user_input.lower() in {"exit", "quit"}:
+            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 __name__ == "__main__":
+    main()