add main.py

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)
+"""