fix: main.py

main.py

@@ -1,7 +1,5 @@
 import os
 import asyncio
-import json
-import httpx
 from pathlib import Path
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_chroma import Chroma
@@ -10,46 +8,49 @@ from langchain.tools import tool
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 
-# --------------------- Configuration ---------------------
+# ----------------- Configuration -----------------
-BASE_DIR = Path(__file__).parent
+# Load OpenRouter API key from .env or environment variable
-DATA_DIR = BASE_DIR / "data"
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-CHROMA_DIR = BASE_DIR / "chroma_faq"
+if not OPENAI_API_KEY:
-META_JSON = BASE_DIR / "course_meta.json"
+    raise RuntimeError("OPENAI_API_KEY not set in environment")
 
-# --------------------- LLM & Embeddings ---------------------
+# ----------------- LLM and Embeddings -----------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
+    api_key=OPENAI_API_KEY,
     temperature=0.0,
 )
 
 embeddings = OpenAIEmbeddings(
     model="text-embedding-3-small",
     base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
+    api_key=OPENAI_API_KEY,
 )
 
-# --------------------- Chroma DB ---------------------
+# ----------------- ChromaDB setup -----------------
+CHROMA_PATH = Path("./chroma_faq")
+CHROMA_COLLECTION = "faq_collection"
 vector_store = Chroma(
-    collection_name="faq_collection",
+    collection_name=CHROMA_COLLECTION,
     embedding_function=embeddings,
-    persist_directory=str(CHROMA_DIR),
+    persist_directory=str(CHROMA_PATH),
 )
 
 # Load markdown files into Chroma if not already loaded
-if not CHROMA_DIR.exists() or not any(CHROMA_DIR.iterdir()):
+if not CHROMA_PATH.exists() or not any(CHROMA_PATH.iterdir()):
+    data_dir = Path("data")
     docs = []
-    for md_file in DATA_DIR.glob("*.md"):
+    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 ---------------------
+# ----------------- Tools -----------------
 @tool
 def search_course_docs(query: str) -> str:
-    """Search the FAQ knowledge base for relevant information."""
+    """Search the local FAQ collection for relevant passages."""
     results = vector_store.similarity_search(query, k=3)
     if not results:
         return "No relevant information found in the course materials."
@@ -57,67 +58,63 @@ def search_course_docs(query: str) -> str:
 
 @tool
 def fetch_course_meta(query: str) -> str:
-    """Fetch course metadata (e.g., schedule) from a local JSON mock."""
+    """Mock MCP-style tool that returns course metadata.
-    if not META_JSON.exists():
+    In production this would be an HTTP call to an MCP server.
-        return "Metadata file not found."
+    Here we return a static JSON-like string based on the query.
-    data = json.loads(META_JSON.read_text(encoding="utf-8"))
+    """
-    # Simple keyword search in the metadata
+    # Simple static mapping for demo purposes
-    matches = [f"{k}: {v}" for k, v in data.items() if query.lower() in k.lower() or query.lower() in str(v).lower()]
+    meta = {
-    return "\n".join(matches) if matches else "No metadata matches the query."
+        "schedule": "Monday 10:00-12:00, Wednesday 14:00-16:00",
+        "instructor": "Dr. Ivanov",
+        "location": "Room 101",
+    }
+    key = query.lower().strip()
+    return meta.get(key, f"No metadata found for '{query}'.")
 
-# --------------------- Backend ---------------------
+# ----------------- Backend -----------------
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
-# --------------------- Agent ---------------------
+# ----------------- Agent -----------------
 agent = create_deep_agent(
     model=llm,
     tools=[search_course_docs, fetch_course_meta],
     backend=backend,
     system_prompt=(
-        "You are a helpful FAQ assistant for the course.\n"
+        "You are a helpful FAQ bot for the course.\n"
-        "When a user asks about course content, use the search_course_docs tool.\n"
+        "When a user asks about course materials, use the search_course_docs tool.\n"
-        "When a user asks about schedule, metadata, or other non‑content info, use fetch_course_meta.\n"
+        "When a user asks about schedule, instructor, or location, use the fetch_course_meta tool.\n"
         "Do not use both tools unless absolutely necessary.\n"
-        "In your final answer, prefix the source with 'source: chroma' or 'source: mcp_meta'."
+        "In your final answer, prefix the response with 'source: chroma' or 'source: mcp_meta' to indicate where the information came from."
     ),
 )
 
-# --------------------- CLI ---------------------
+# ----------------- CLI -----------------
-SAMPLE_QUESTIONS = [
+PRESET_QUESTIONS = [
-    "What topics are covered in the first lecture?",  # should hit chroma
+    "What topics are covered in the first lecture?",
-    "Explain the concept of tokenization in NLP.",  # chroma
+    "When is the next class?",
-    "When is the next class scheduled?",  # meta
+    "Who is the instructor?",
 ]
 
-async def run_interactive():
+async def run_cli():
-    print("Welcome to the Course FAQ Bot! Type 'exit' to quit.")
+    print("Welcome to the Course FAQ Bot!\n")
+    for i, q in enumerate(PRESET_QUESTIONS, 1):
+        print(f"{i}. {q}")
+    print("\nEnter your own question or type 'exit' to quit.")
     while True:
-        user_input = input("\nYou: ")
+        user_input = input("\n> ")
         if user_input.lower() in {"exit", "quit"}:
+            print("Goodbye!")
             break
         result = await agent.ainvoke(
-            {"messages": [{"role": "user", "content": user_input}]},
+            {"messages": ["HumanMessage(content=\"{}\")".format(user_input)]},
-            {"configurable": {"thread_id": "interactive-session"}},
+            {"configurable": {"thread_id": "session-1"}},
         )
-        print("\nAssistant:", result["messages"][-1]["content"])
+        # The agent returns a dict with 'messages'; take the last one
-
+        content = result["messages"][-1].content
-async def run_samples():
+        print(content)
-    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():
-    # Run sample questions first
-    await run_samples()
-    # Then enter interactive mode
-    await run_interactive()
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    asyncio.run(run_cli())