add: main.py

main.py

@@ -1,73 +1,22 @@
 import os
 import asyncio
 import json
-from pathlib import Path
-from typing import List
-
 import httpx
-from langchain_openai import ChatOpenAI
+from pathlib import Path
-from langchain_core.messages import HumanMessage
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_chroma import Chroma
+from langchain_core.documents import Document
 from langchain.tools import tool
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-from langchain_chroma import Chroma
-from langchain_ollama import OllamaEmbeddings
 
-# ---------------------
+# --------------------- Configuration ---------------------
-# 1. Chroma DB helpers
+BASE_DIR = Path(__file__).parent
-# ---------------------
+DATA_DIR = BASE_DIR / "data"
-CHROMA_PATH = Path("./chroma_faq")
+CHROMA_DIR = BASE_DIR / "chroma_faq"
-DATA_DIR = Path("./data")
+META_JSON = BASE_DIR / "course_meta.json"
 
-
+# --------------------- LLM & Embeddings ---------------------
-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",
@@ -75,56 +24,100 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
+embeddings = OpenAIEmbeddings(
+    model="text-embedding-3-small",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENAI_API_KEY"),
+)
+
+# --------------------- Chroma DB ---------------------
+vector_store = Chroma(
+    collection_name="faq_collection",
+    embedding_function=embeddings,
+    persist_directory=str(CHROMA_DIR),
+)
+
+# Load markdown files into Chroma if not already loaded
+if not CHROMA_DIR.exists() or not any(CHROMA_DIR.iterdir()):
+    docs = []
+    for md_file in DATA_DIR.glob("*.md"):
+        text = md_file.read_text(encoding="utf-8")
+        docs.append(Document(page_content=text, metadata={"source": md_file.name}))
+    vector_store.add_documents(docs)
+    vector_store.persist()
+
+# --------------------- Tools ---------------------
+@tool
+def search_course_docs(query: str) -> str:
+    """Search the FAQ knowledge base for relevant information."""
+    results = vector_store.similarity_search(query, k=3)
+    if not results:
+        return "No relevant information found in the course materials."
+    return "\n---\n".join(f"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc in results)
+
+@tool
+def fetch_course_meta(query: str) -> str:
+    """Fetch course metadata (e.g., schedule) from a local JSON mock."""
+    if not META_JSON.exists():
+        return "Metadata file not found."
+    data = json.loads(META_JSON.read_text(encoding="utf-8"))
+    # Simple keyword search in the metadata
+    matches = [f"{k}: {v}" for k, v in data.items() if query.lower() in k.lower() or query.lower() in str(v).lower()]
+    return "\n".join(matches) if matches else "No metadata matches the query."
+
+# --------------------- Backend ---------------------
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
-SYSTEM_PROMPT = (
+# --------------------- Agent ---------------------
-    "You are a helpful FAQ assistant for the course.\n"
-    "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`."
-)
-
 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 assistant for the course.\n"
+        "When a user asks about course content, use the search_course_docs tool.\n"
+        "When a user asks about schedule, metadata, or other non‑content info, use fetch_course_meta.\n"
+        "Do not use both tools unless absolutely necessary.\n"
+        "In your final answer, prefix the source with 'source: chroma' or 'source: mcp_meta'."
+    ),
 )
 
-# ---------------------
+# --------------------- CLI ---------------------
-# 4. CLI
+SAMPLE_QUESTIONS = [
-# ---------------------
+    "What topics are covered in the first lecture?",  # should hit chroma
-PRESET_QUESTIONS = [
+    "Explain the concept of tokenization in NLP.",  # chroma
-    "What is the deadline for the final project?",  # FAQ
+    "When is the next class scheduled?",  # meta
-    "When does the next lecture on deep learning start?",  # metadata
-    "Explain the concept of attention mechanism.",  # FAQ
 ]
 
-async def run_agent(question: str) -> str:
+async def run_interactive():
-    result = await agent.ainvoke(
+    print("Welcome to the Course FAQ Bot! Type 'exit' to quit.")
-        {"messages": [HumanMessage(content=question)]},
+    while True:
-        {"configurable": {"thread_id": "session-1"}},
+        user_input = input("\nYou: ")
-    )
+        if user_input.lower() in {"exit", "quit"}:
-    return result["messages"][-1].content
+            break
+        result = await agent.ainvoke(
+            {"messages": [{"role": "user", "content": user_input}]},
+            {"configurable": {"thread_id": "interactive-session"}},
+        )
+        print("\nAssistant:", result["messages"][-1]["content"])
+
+async def run_samples():
+    for q in SAMPLE_QUESTIONS:
+        print("\nQuestion:", q)
+        result = await agent.ainvoke(
+            {"messages": [{"role": "user", "content": q}]},
+            {"configurable": {"thread_id": "sample-session"}},
+        )
+        print("Answer:", result["messages"][-1]["content"])
 
 async def main():
-    print("--- FAQ Bot Demo ---\n")
+    # Run sample questions first
-    for i, q in enumerate(PRESET_QUESTIONS, 1):
+    await run_samples()
-        print(f"Q{i}: {q}")
+    # Then enter interactive mode
-        ans = await run_agent(q)
+    await run_interactive()
-        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
-        ans = await run_agent(user_q)
-        print(ans)
 
 if __name__ == "__main__":
     asyncio.run(main())