diff --git a/main.py b/main.py index 5dc42dc..a80d378 100644 --- a/main.py +++ b/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()