From 0e6b2ad4c080fad2a92f3acaebaadcca7203cd31 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=AD=D0=BC=D0=B8=D0=BB=D1=8C=20=D0=90=D0=BC=D0=B8=D1=80?= =?UTF-8?q?=D0=BE=D0=B2?= Date: Fri, 5 Jun 2026 11:18:11 +0000 Subject: [PATCH] add main.py --- main.py | 139 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..5dc42dc --- /dev/null +++ b/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()