From ac3105de434ce7312051409624065053745aa12b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=9A=D1=83=D1=82?= =?UTF-8?q?=D0=BB=D0=B0=D1=85=D0=BC=D0=B5=D1=82=D0=BE=D0=B2?= Date: Thu, 4 Jun 2026 16:00:12 +0000 Subject: [PATCH] add main.py --- main.py | 185 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 185 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..9e6a875 --- /dev/null +++ b/main.py @@ -0,0 +1,185 @@ +""" +Main entry point for the FAQ‑bot assignment. + +The bot answers questions about course materials stored in a local ChromaDB +vector store and, when necessary, calls an external MCP‑style tool that +returns metadata such as the course schedule. The implementation uses +LangChain v1+ (``langchain>=1.0.0``) together with ``langchain-chroma`` +and ``langchain-ollama`` for embeddings. + +The file contains: +* Utility functions to load markdown files into ChromaDB. +* Two tools – one that searches the vector store and another that mocks an + MCP call. +* A simple LangChain agent that decides which tool to use based on the + content of the user query. +* Three example queries executed when the script is run as ``__main__``. +""" + +from __future__ import annotations + +import os +import json +import pathlib +from typing import List, Dict + +# --------------------------------------------------------------------------- +# Dependencies – all are declared in requirements.txt +# --------------------------------------------------------------------------- +from langchain_ollama import ChatOllama +from langchain_chroma import Chroma +from langchain_core.messages import HumanMessage +from langchain.agents import create_agent +from langchain.tools import tool +from langchain.embeddings.ollama import OllamaEmbeddings + +# --------------------------------------------------------------------------- +# Configuration constants +# --------------------------------------------------------------------------- +DATA_DIR = pathlib.Path("data") # directory with .md files +CHROMA_PATH = pathlib.Path("./chroma_faq") # persistent store location +MCP_ENDPOINT = "http://localhost:8000/course_meta" # mock endpoint + +# --------------------------------------------------------------------------- +# Helper – load markdown files into ChromaDB +# --------------------------------------------------------------------------- +@tool(name="load_faq_to_chroma", description="Load all .md files from data/ into a persistent Chroma vector store.") +def load_faq_to_chroma() -> str: + """Create or update the Chroma collection with local markdown files. + + The function is idempotent – it will overwrite existing documents + with the same IDs. It returns a short JSON string describing how many + documents were added. + """ + # Ensure data directory exists + if not DATA_DIR.exists(): + return json.dumps({"error": f"{DATA_DIR} does not exist."}) + + # Gather all markdown files + md_files = list(DATA_DIR.glob("*.md")) + if not md_files: + return json.dumps({"error": "No .md files found in data/"}) + + # Prepare documents for Chroma + docs: List[Dict[str, str]] = [] + for file_path in md_files: + content = file_path.read_text(encoding="utf-8") + docs.append({"id": file_path.stem, "content": content}) + + # Create embeddings and store + embedding = OllamaEmbeddings(model="nomic-embed-text") + chroma = Chroma.from_documents( + documents=[doc["content"] for doc in docs], + ids=[doc["id"] for doc in docs], + embedding=embedding, + persist_directory=str(CHROMA_PATH), + ) + + return json.dumps({"status": "loaded", "documents": len(docs)}) + +# --------------------------------------------------------------------------- +# Tool – search the vector store +# --------------------------------------------------------------------------- +@tool(name="search_course_docs", description="Search local course documents for a query. Returns top k results.") +def search_course_docs(query: str, k: int = 3) -> str: + """Return the most relevant chunks from the Chroma store. + + The function loads the persistent collection and performs a similarity + search. Results are returned as a numbered list of snippets. + """ + if not CHROMA_PATH.exists(): + return json.dumps({"error": f"Chroma store {CHROMA_PATH} not found."}) + + embedding = OllamaEmbeddings(model="nomic-embed-text") + chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=embedding) + docs = chroma.similarity_search(query, k=k) + if not docs: + return json.dumps({"result": "No relevant documents found."}) + + snippets = [f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)] + return json.dumps({"source": "chroma", "snippets": snippets}) + +# --------------------------------------------------------------------------- +# Tool – mock MCP call (could be replaced with a real HTTP request) +# --------------------------------------------------------------------------- +@tool(name="fetch_course_meta", description="Retrieve course metadata such as schedule. Returns JSON.") +def fetch_course_meta(query: str) -> str: + """Simulate an external MCP tool by returning static data. + + In a real deployment this would perform an HTTP GET to the MCP server. + For the purposes of the assignment we simply return a hard‑coded JSON + payload that mimics what an MCP endpoint might provide. + """ + # Static mock data – in practice replace with httpx.get(MCP_ENDPOINT) + mock_data = { + "schedule": [ + {"day": "Monday", "time": "10:00-12:00", "topic": "Introduction"}, + {"day": "Wednesday", "time": "14:00-16:00", "topic": "Advanced Topics"}, + ], + "instructor": "Prof. Alexei Petrov", + } + return json.dumps({"source": "mcp_meta", "data": mock_data}) + +# --------------------------------------------------------------------------- +# Agent – decides which tool to use based on the query content +# --------------------------------------------------------------------------- +# System prompt guides the agent not to call both tools unnecessarily. +SYSTEM_PROMPT = ( + "You are an assistant that answers questions about course materials." + " If the question is about lecture notes or FAQs, use the search_course_docs tool." + " If the question asks for schedule, instructor info, or other metadata, use fetch_course_meta." +) + +# Create the agent once – it will be reused in examples. +llm = ChatOllama(model="llama3", temperature=0.0) +agent = create_agent( + llm=llm, + tools=[search_course_docs, fetch_course_meta], + system_prompt=SYSTEM_PROMPT, +) + +# --------------------------------------------------------------------------- +# Example usage – three predefined queries and an interactive prompt +# --------------------------------------------------------------------------- +EXAMPLES: List[str] = [ + "What topics are covered in the first lecture?", + "When is the next class scheduled?", + "Can you give me a summary of the course material?", +] + +def run_example(query: str) -> None: + print(f"\nUser: {query}") + result = agent.invoke( + {"messages": [HumanMessage(content=query)]}, + {"configurable": {"thread_id": "example-session"}}, + ) + # The agent returns a list of messages – the last one is the answer. + answer = result["messages"][-1].content + print(f"Assistant: {answer}") + +if __name__ == "__main__": + # Ensure data directory exists for demo purposes + if not DATA_DIR.exists(): + DATA_DIR.mkdir(parents=True, exist_ok=True) + # Create a tiny FAQ file as an example. + (DATA_DIR / "faq.md").write_text( + "# Course FAQ\n" \ + "## What is the course about?\n" \ + "The course covers advanced topics in AI and machine learning." + ) + # Load data into Chroma (idempotent) + print(load_faq_to_chroma()) + + for q in EXAMPLES: + run_example(q) + + # Interactive mode – optional + while True: + try: + user_input = input("\nAsk a question (or 'quit'): ") + except EOFError: + break + if user_input.lower() in {"quit", "exit"}: + break + run_example(user_input) +"""