fix: main.py

main.py

@@ -1,120 +1,144 @@
 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 -----------------
+# ---------------------------------------------------------------------------
-# Load OpenRouter API key from .env or environment variable
+# 1. Настройка LLM и Embeddings (OpenRouter)
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+# ---------------------------------------------------------------------------
-if not OPENAI_API_KEY:
-    raise RuntimeError("OPENAI_API_KEY not set in environment")
-
-# ----------------- LLM and Embeddings -----------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
-    api_key=OPENAI_API_KEY,
+    api_key=os.getenv("OPENAI_API_KEY"),
     temperature=0.0,
 )
 
 embeddings = OpenAIEmbeddings(
     model="text-embedding-3-small",
     base_url="https://openrouter.ai/api/v1",
-    api_key=OPENAI_API_KEY,
+    api_key=os.getenv("OPENAI_API_KEY"),
 )
 
-# ----------------- ChromaDB setup -----------------
+# ---------------------------------------------------------------------------
+# 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),
 )
 
-# Load markdown files into Chroma if not already loaded
+# Если коллекция пуста, загрузим данные из data/*.md
-if not CHROMA_PATH.exists() or not any(CHROMA_PATH.iterdir()):
+if not vector_store.get_collection().list_documents():
-    data_dir = Path("data")
+    md_files = list(Path("data").glob("*.md"))
-    docs = []
+    docs: List[Document] = []
-    for md_file in data_dir.glob("*.md"):
+    for f in md_files:
-        text = md_file.read_text(encoding="utf-8")
+        text = f.read_text(encoding="utf-8")
-        docs.append(Document(page_content=text, metadata={"source": md_file.name}))
+        docs.append(Document(page_content=text, metadata={"source": f.name}))
     vector_store.add_documents(docs)
     vector_store.persist()
 
-# ----------------- Tools -----------------
 @tool
-def search_course_docs(query: str) -> str:
+def search_course_docs(query: str, k: int = 3) -> str:
     """Search the local FAQ collection for relevant passages."""
-    results = vector_store.similarity_search(query, k=3)
+    results = vector_store.similarity_search(query, k=k)
     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)
+    return "\n\n---\n\n".join(f"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc in results)
+
+# ---------------------------------------------------------------------------
+# 3. MCP‑style tool (mocked via local JSON file)
+# ---------------------------------------------------------------------------
+META_JSON = Path("course_meta.json")
+if not META_JSON.exists():
+    # Создаём простую статическую мета‑информацию
+    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:
-    """Mock MCP-style tool that returns course metadata.
+    """Return course metadata that matches the query.
-    In production this would be an HTTP call to an MCP server.
+    The function simply looks for the query string in the keys of the JSON.
-    Here we return a static JSON-like string based on the query.
     """
-    # Simple static mapping for demo purposes
+    data = json.loads(META_JSON.read_text())
-    meta = {
+    for key, value in data.items():
-        "schedule": "Monday 10:00-12:00, Wednesday 14:00-16:00",
+        if query.lower() in key.lower():
-        "instructor": "Dr. Ivanov",
+            return json.dumps({key: value}, indent=2)
-        "location": "Room 101",
+    return "No metadata found for the given query."
-    }
-    key = query.lower().strip()
-    return meta.get(key, f"No metadata found for '{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 the course.\n"
+        "You are a helpful FAQ bot for a course.\n"
-        "When a user asks about course materials, use the search_course_docs tool.\n"
+        "If the question is about course content, use the search_course_docs tool.\n"
-        "When a user asks about schedule, instructor, or location, use the fetch_course_meta tool.\n"
+        "If the question is about schedule, instructor, or other metadata, use fetch_course_meta.\n"
-        "Do not use both tools unless absolutely necessary.\n"
+        "Do not call both tools unless absolutely necessary.\n"
-        "In your final answer, prefix the response with 'source: chroma' or 'source: mcp_meta' to indicate where the information came from."
+        "In your final answer, prepend 'source: chroma' or 'source: mcp_meta' to indicate where the answer came from."
     ),
 )
 
-# ----------------- CLI -----------------
+# ---------------------------------------------------------------------------
+# 5. CLI
+# ---------------------------------------------------------------------------
 PRESET_QUESTIONS = [
-    "What topics are covered in the first lecture?",
+    "What topics are covered in Lecture 1?",  # should hit chroma
-    "When is the next class?",
+    "Who is the instructor for this course?",  # should hit mcp_meta
-    "Who is the instructor?",
+    "Explain the concept of polymorphism."
 ]
 
-async def run_cli():
+async def run_agent(question: str) -> str:
-    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("\n> ")
-        if user_input.lower() in {"exit", "quit"}:
-            print("Goodbye!")
-            break
     result = await agent.ainvoke(
-            {"messages": ["HumanMessage(content=\"{}\")".format(user_input)]},
+        {"messages": [HumanMessage(content=question)]},
         {"configurable": {"thread_id": "session-1"}},
     )
-        # The agent returns a dict with 'messages'; take the last one
+    return result["messages"][-1].content
-        content = result["messages"][-1].content
+
-        print(content)
+async def main():
+    print("\n--- 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 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(run_cli())
+    asyncio.run(main())