add main.py

main.py

@@ -1,5 +1,10 @@
-import asyncio, os
+import os
+import asyncio
+import json
 from pathlib import Path
+from typing import List
+
+import httpx
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
 from langchain.tools import tool
@@ -7,10 +12,62 @@ from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 from langchain_chroma import Chroma
 from langchain_ollama import OllamaEmbeddings
-import httpx
-import json
 
-# ---------- LLM ----------
+# ---------------------
+# 1. Chroma DB helpers
+# ---------------------
+CHROMA_PATH = Path("./chroma_faq")
+DATA_DIR = Path("./data")
+
+
+def load_faq_to_chroma() -> None:
+    """Load all .md files from data/ into a persistent Chroma store."""
+    if CHROMA_PATH.exists():
+        # already loaded
+        return
+    CHROMA_PATH.mkdir(parents=True, exist_ok=True)
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    db = Chroma.from_folder(str(DATA_DIR), embedding=embeddings, persist_directory=str(CHROMA_PATH))
+    db.persist()
+
+@tool
+def search_course_docs(query: str, k: int = 3) -> str:
+    """Search local course FAQ documents in Chroma DB."""
+    load_faq_to_chroma()
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    db = Chroma(persist_directory=str(CHROMA_PATH), embedding=embeddings)
+    docs = db.similarity_search(query, k=k)
+    if not docs:
+        return "No relevant FAQ found."
+    return "\n\n---\n\n".join(doc.page_content for doc in docs)
+
+# ---------------------
+# 2. MCP‑style tool
+# ---------------------
+# For demo we use a local JSON file served by python -m http.server
+# The file is located at ./meta/course_meta.json
+
+META_URL = "http://localhost:8000/course_meta.json"
+
+@tool
+def fetch_course_meta(query: str) -> str:
+    """Fetch course metadata (e.g., schedule) from a mock MCP server."""
+    try:
+        response = httpx.get(META_URL, timeout=5.0)
+        response.raise_for_status()
+        data = response.json()
+    except Exception as e:
+        return f"Error fetching metadata: {e}"
+    # Simple keyword search in the JSON
+    results: List[str] = []
+    for key, value in data.items():
+        if query.lower() in key.lower() or query.lower() in str(value).lower():
+            results.append(f"{key}: {value}")
+    return "\n".join(results) if results else "No metadata matches your query."
+
+# ---------------------
+# 3. Agent setup
+# ---------------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -18,97 +75,56 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ---------- Backend ----------
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
-# ---------- Chroma DB ----------
+SYSTEM_PROMPT = (
-CHROMA_PATH = Path("./chroma_faq")
+    "You are a helpful FAQ assistant for the course.\n"
-CHROMA_PATH.mkdir(exist_ok=True)
+    "When a user asks a question, first decide whether the answer is best found in the local FAQ documents or in the course metadata.\n"
+    "If the answer is in the FAQ, use the tool `search_course_docs`.\n"
+    "If the answer requires schedule or other metadata, use the tool `fetch_course_meta`.\n"
+    "Do not call both tools unless absolutely necessary.\n"
+    "In your final response, prepend the source: `source: chroma` or `source: mcp_meta`."
+)
 
-embeddings = OllamaEmbeddings(model="nomic-embed-text")
-
-# Load or create vector store
-if CHROMA_PATH.exists() and any(CHROMA_PATH.iterdir()):
-    chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=embeddings)
-else:
-    # Load markdown files
-    docs = []
-    for md_file in Path("data").glob("*.md"):
-        text = md_file.read_text(encoding="utf-8")
-        docs.append(text)
-    chroma = Chroma.from_texts(docs, embedding=embeddings, persist_directory=str(CHROMA_PATH))
-    chroma.persist()
-
-@tool
-def search_course_docs(query: str, k: int = 3) -> str:
-    """Search local course documents in Chroma."""
-    results = chroma.similarity_search(query, k=k)
-    return "\n---\n".join(doc.page_content for doc in results) if results else "No relevant docs found."
-
-# ---------- MCP‑style tool ----------
-# For demo we use a static JSON file. In production this would be an HTTP call.
-META_JSON = Path("meta.json")
-if not META_JSON.exists():
-    # Create a simple mock meta file
-    META_JSON.write_text(json.dumps({
-        "schedule": "Mon 10-12, Wed 14-16, Fri 9-11",
-        "instructor": "Dr. Smith",
-        "location": "Room 101"
-    }))
-
-@tool
-def fetch_course_meta(query: str) -> str:
-    """Return course metadata matching the query keyword."""
-    data = json.loads(META_JSON.read_text())
-    # Simple keyword search in values
-    for key, value in data.items():
-        if query.lower() in key.lower() or query.lower() in str(value).lower():
-            return f"{key}: {value}"
-    return "No metadata found for the query."
-
-# ---------- Agent ----------
 agent = create_deep_agent(
     model=llm,
     tools=[search_course_docs, fetch_course_meta],
     backend=backend,
-    system_prompt=(
+    system_prompt=SYSTEM_PROMPT,
-        "You are a helpful FAQ bot for the course.\n"
-        "If the question is about course content, use search_course_docs.\n"
-        "If the question is about schedule, instructor, or location, use fetch_course_meta.\n"
-        "Do not call both tools unless necessary.\n"
-        "In your answer, prefix the source with 'source: chroma' or 'source: mcp_meta'."
-    ),
 )
 
-# ---------- CLI ----------
+# ---------------------
+# 4. CLI
+# ---------------------
 PRESET_QUESTIONS = [
-    "What is the main topic of the first lecture?",
+    "What is the deadline for the final project?",  # FAQ
-    "How can I access the lecture slides?",
+    "When does the next lecture on deep learning start?",  # metadata
-    "When is the next class?"
+    "Explain the concept of attention mechanism.",  # FAQ
 ]
 
-async def run_cli():
+async def run_agent(question: str) -> str:
-    print("--- FAQ Bot CLI ---")
-    for i, q in enumerate(PRESET_QUESTIONS, 1):
-        print(f"\nPreset {i}: {q}")
     result = await agent.ainvoke(
-            {"messages": [HumanMessage(content=q)]},
+        {"messages": [HumanMessage(content=question)]},
-            {"configurable": {"thread_id": f"preset-{i}"}},
+        {"configurable": {"thread_id": "session-1"}},
     )
-        print(result["messages"][-1].content)
+    return result["messages"][-1].content
-    print("\nEnter your own question (or 'exit'): ")
+
+async def main():
+    print("--- FAQ Bot Demo ---\n")
+    for i, q in enumerate(PRESET_QUESTIONS, 1):
+        print(f"Q{i}: {q}")
+        ans = await run_agent(q)
+        print(f"A{i}: {ans}\n")
+    print("Enter your own question (or 'exit' to quit):")
     while True:
         user_q = input("> ")
         if user_q.lower() in {"exit", "quit"}:
             break
-        result = await agent.ainvoke(
+        ans = await run_agent(user_q)
-            {"messages": [HumanMessage(content=user_q)]},
+        print(ans)
-            {"configurable": {"thread_id": "interactive"}},
-        )
-        print(result["messages"][-1].content)
 
 if __name__ == "__main__":
-    asyncio.run(run_cli())
+    asyncio.run(main())