fix: main.py

main.py

@@ -1,20 +1,18 @@
 import os
 import asyncio
-import json
 from pathlib import Path
-from typing import List
-
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_chroma import Chroma
 from langchain_core.documents import Document
-from langchain_core.messages import HumanMessage
 from langchain.tools import tool
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 
-# ---------------------------------------------------------------------------
+# ----------------- Configuration -----------------
-# 1. Настройка LLM и Embeddings (OpenRouter)
+# Load API key from .env or environment variable
-# ---------------------------------------------------------------------------
+os.environ.setdefault("OPENAI_API_KEY", os.getenv("OPENAI_API_KEY", ""))
+
+# LLM via OpenRouter
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -22,123 +20,105 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
+# Embeddings for Chroma
 embeddings = OpenAIEmbeddings(
     model="text-embedding-3-small",
     base_url="https://openrouter.ai/api/v1",
     api_key=os.getenv("OPENAI_API_KEY"),
 )
 
-# ---------------------------------------------------------------------------
+# ----------------- Chroma DB -----------------
-# 2. ChromaDB: загрузка .md файлов и поиск
-# ---------------------------------------------------------------------------
 CHROMA_PATH = Path("./chroma_faq")
 CHROMA_COLLECTION = "faq_collection"
-
 vector_store = Chroma(
     collection_name=CHROMA_COLLECTION,
     embedding_function=embeddings,
     persist_directory=str(CHROMA_PATH),
 )
 
-# Если коллекция пуста, загрузим данные из data/*.md
+# Load markdown files into Chroma (idempotent)
-if not vector_store.get_collection().list_documents():
+DATA_DIR = Path("./data")
-    md_files = list(Path("data").glob("*.md"))
+if not CHROMA_PATH.exists() or not list(CHROMA_PATH.iterdir()):
-    docs: List[Document] = []
+    docs = []
-    for f in md_files:
+    for md_file in DATA_DIR.glob("*.md"):
-        text = f.read_text(encoding="utf-8")
+        text = md_file.read_text(encoding="utf-8")
-        docs.append(Document(page_content=text, metadata={"source": f.name}))
+        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, k: int = 3) -> str:
+def search_course_docs(query: str) -> str:
     """Search the local FAQ collection for relevant passages."""
-    results = vector_store.similarity_search(query, k=k)
+    results = vector_store.similarity_search(query, k=3)
     if not results:
         return "No relevant information found in the course materials."
-    return "\n\n---\n\n".join(f"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc in results)
+    return "\n---\n".join(f"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc in results)
 
-# ---------------------------------------------------------------------------
+# Mock MCP tool – static JSON data
-# 3. MCP‑style tool (mocked via local JSON file)
+COURSE_META = {
-# ---------------------------------------------------------------------------
+    "schedule": "Monday 10:00-12:00, Wednesday 14:00-16:00",
-META_JSON = Path("course_meta.json")
+    "instructor": "Dr. Ivanov",
-if not META_JSON.exists():
+    "location": "Room 101",
-    # Создаём простую статическую мета‑информацию
+}
-    META_JSON.write_text(json.dumps({
-        "schedule": {
-            "Monday": "Lecture 1",
-            "Wednesday": "Lecture 2",
-            "Friday": "Lab"
-        },
-        "instructor": "Dr. Example"
-    }, indent=2))
 
 @tool
 def fetch_course_meta(query: str) -> str:
-    """Return course metadata that matches the query.
+    """Return course metadata matching the query keyword.
-    The function simply looks for the query string in the keys of the JSON.
+    For example, query="schedule" returns the schedule string.
     """
-    data = json.loads(META_JSON.read_text())
+    key = query.lower().strip()
-    for key, value in data.items():
+    return COURSE_META.get(key, f"No metadata found for '{query}'.")
-        if query.lower() in key.lower():
-            return json.dumps({key: value}, indent=2)
-    return "No metadata found for the given query."
 
-# ---------------------------------------------------------------------------
+# ----------------- Backend -----------------
-# 4. DeepAgent с маршрутизацией
-# ---------------------------------------------------------------------------
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
+# ----------------- Agent -----------------
 agent = create_deep_agent(
     model=llm,
     tools=[search_course_docs, fetch_course_meta],
     backend=backend,
     system_prompt=(
-        "You are a helpful FAQ bot for a course.\n"
+        "You are a helpful FAQ assistant for the course.\n"
-        "If the question is about course content, use the search_course_docs tool.\n"
+        "When answering a question, first decide whether the answer comes from the course materials (use search_course_docs)\n"
-        "If the question is about schedule, instructor, or other metadata, use fetch_course_meta.\n"
+        "or from course metadata (use fetch_course_meta).\n"
         "Do not call both tools unless absolutely necessary.\n"
-        "In your final answer, prepend 'source: chroma' or 'source: mcp_meta' to indicate where the answer came from."
+        "In your final answer, prepend 'source: chroma' or 'source: mcp_meta' to indicate the origin."
     ),
 )
 
-# ---------------------------------------------------------------------------
+# ----------------- CLI -----------------
-# 5. CLI
-# ---------------------------------------------------------------------------
 PRESET_QUESTIONS = [
-    "What topics are covered in Lecture 1?",  # should hit chroma
+    "What topics are covered in the first lecture?",  # should hit chroma
     "Who is the instructor for this course?",          # should hit mcp_meta
-    "Explain the concept of polymorphism."
+    "When is the next class?",                        # should hit chroma (or meta if schedule)
 ]
 
-async def run_agent(question: str) -> str:
+async def run_question(question: str):
     result = await agent.ainvoke(
-        {"messages": [HumanMessage(content=question)]},
+        {"messages": ["HumanMessage(content=\"{}\")".format(question)]},
         {"configurable": {"thread_id": "session-1"}},
     )
-    return result["messages"][-1].content
+    # The agent returns a dict with 'messages'; extract last content
+    content = result["messages"][-1].content
+    print(f"\nQ: {question}\nA: {content}\n")
+
+async def interactive():
+    print("Enter a question (or 'exit' to quit):")
+    while True:
+        q = input("> ")
+        if q.lower() in {"exit", "quit"}:
+            break
+        await run_question(q)
 
 async def main():
-    print("\n--- FAQ Bot Demo ---\n")
+    print("Running preset questions...")
-    for i, q in enumerate(PRESET_QUESTIONS, 1):
+    for q in PRESET_QUESTIONS:
-        print(f"Q{i}: {q}")
+        await run_question(q)
-        ans = await run_agent(q)
+    await interactive()
-        print(f"A{i}: {ans}\n")
-
-    print("Enter your own question (or press Ctrl+C to exit):")
-    while True:
-        try:
-            user_q = input("> ")
-            if not user_q.strip():
-                continue
-            ans = await run_agent(user_q)
-            print(ans)
-        except KeyboardInterrupt:
-            print("\nExiting.")
-            break
 
 if __name__ == "__main__":
     asyncio.run(main())