fix(needs_fixes): 1 исправлений, 0 отстояно — main.py

main.py

@@ -1,19 +1,20 @@
 import os
 import asyncio
 import json
-import httpx
 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
-from langchain_core.messages import HumanMessage
 
-# ---------------------------
+# ---------------------------------------------------------------------------
 # Configuration
-# ---------------------------
+# ---------------------------------------------------------------------------
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 if not OPENAI_API_KEY:
     raise RuntimeError("OPENAI_API_KEY not set in environment")
@@ -26,93 +27,111 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# Embeddings for Chroma
+# Embeddings via OpenRouter
 embeddings = OpenAIEmbeddings(
     model="text-embedding-3-small",
     base_url="https://openrouter.ai/api/v1",
     api_key=OPENAI_API_KEY,
 )
 
-# ---------------------------
+# ---------------------------------------------------------------------------
-# Chroma DB utilities
+# Chroma persistence
-# ---------------------------
+# ---------------------------------------------------------------------------
 CHROMA_DIR = Path("./chroma_faq")
 CHROMA_DIR.mkdir(parents=True, exist_ok=True)
 vector_store = Chroma(
     collection_name="faq_collection",
-    persist_directory=str(CHROMA_DIR),
     embedding_function=embeddings,
+    persist_directory=str(CHROMA_DIR),
 )
 
-
+# ---------------------------------------------------------------------------
-def load_faq_to_chroma(md_dir: str = "data"):
+# Utility: load markdown files into Chroma
-    """Load all .md files from md_dir into ChromaDB.
+# ---------------------------------------------------------------------------
-    Each file is split into chunks and added to the vector store.
+@tool
+def load_faq_to_chroma() -> str:
+    """Load all .md files from data/ into the Chroma vector store.
+    This tool is idempotent – it will overwrite existing collection.
     """
-    md_path = Path(md_dir)
+    data_dir = Path("data")
-    if not md_path.exists():
+    if not data_dir.exists():
-        raise FileNotFoundError(f"Markdown directory {md_dir} not found")
+        return "Data directory not found."
-    for md_file in md_path.glob("*.md"):
+    docs: List[Document] = []
+    for md_file in data_dir.glob("*.md"):
         text = md_file.read_text(encoding="utf-8")
-        # Simple chunking: split by double newlines
+        docs.append(Document(page_content=text, metadata={"source": md_file.name}))
-        chunks = [c.strip() for c in text.split("\n\n") if c.strip()]
+    if docs:
-        docs = [Document(page_content=c, metadata={"source": md_file.name}) for c in chunks]
+        vector_store.delete_collection()
         vector_store.add_documents(docs)
-    vector_store.persist()
+        vector_store.persist()
+        return f"Loaded {len(docs)} documents into Chroma."
+    return "No markdown files found."
 
-# ---------------------------
+# ---------------------------------------------------------------------------
-# Tools
+# Tool: search knowledge base
-# ---------------------------
+# ---------------------------------------------------------------------------
 @tool
 def search_course_docs(query: str, k: int = 3) -> str:
-    """Search the local FAQ ChromaDB for relevant passages."""
+    """Search the FAQ collection for relevant passages.
+    Returns a string with the top k passages and a source tag.
+    """
     docs = vector_store.similarity_search(query, k=k)
     if not docs:
-        return "No relevant information found in the course materials."
+        return "No relevant information found in the FAQ."
-    return "\n\n---\n\n".join(f"{d.metadata.get('source', 'unknown')}\n{d.page_content}" for d in docs)
+    results = "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs))
+    return f"source: chroma\n{results}"
+
+# ---------------------------------------------------------------------------
+# MCP-style tool: fetch course metadata
+# ---------------------------------------------------------------------------
+# For simplicity we use a static JSON file in the repo. In production this
+# would be an HTTP call to an MCP server.
+METADATA_FILE = Path("course_meta.json")
 
 @tool
 def fetch_course_meta(query: str) -> str:
-    """Mock MCP-style tool that fetches course metadata from a local JSON file.
+    """Return metadata that matches the query.
-    In production this would be an HTTP call to an MCP server.
+    The function performs a simple keyword search in the static JSON.
     """
-    # For simplicity, we use a local JSON file. In a real scenario, replace with httpx.get.
+    if not METADATA_FILE.exists():
-    meta_path = Path("course_meta.json")
+        return "Metadata file not found."
-    if not meta_path.exists():
+    data = json.loads(METADATA_FILE.read_text(encoding="utf-8"))
-        return "Course metadata not available."
+    matches = [item for item in data if query.lower() in item.get("title", "").lower()]
-    data = json.loads(meta_path.read_text(encoding="utf-8"))
+    if not matches:
-    # Very naive search: return any entry where query is a substring of title or description
+        return "No metadata matches the query."
-    results = [f"{item['title']}: {item['description']}" for item in data if query.lower() in item.get('title', '').lower() or query.lower() in item.get('description', '').lower()]
+    return f"source: mcp_meta\n" + json.dumps(matches, indent=2)
-    return "\n".join(results) if results else "No matching metadata found."
 
-# ---------------------------
+# ---------------------------------------------------------------------------
-# Agent setup
+# Backend setup
-# ---------------------------
+# ---------------------------------------------------------------------------
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
+# ---------------------------------------------------------------------------
+# Agent definition
+# ---------------------------------------------------------------------------
 agent = create_deep_agent(
     model=llm,
-    tools=[search_course_docs, fetch_course_meta],
+    tools=[load_faq_to_chroma, 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. "
-        "When a user asks about course content, use the search_course_docs tool.\n"
+        "Use the search_course_docs tool for questions about lecture materials. "
-        "When a user asks about schedule, metadata, or other non-content info, use fetch_course_meta.\n"
+        "Use fetch_course_meta for questions about schedule or metadata. "
-        "Do not call both tools unless absolutely necessary.\n"
+        "Do not call both tools unless necessary. "
-        "In your final answer, prepend 'source: chroma' or 'source: mcp_meta' to indicate the origin."
+        "Always prefix your answer with the source tag (chroma or mcp_meta)."
     ),
 )
 
-# ---------------------------
+# ---------------------------------------------------------------------------
-# CLI
+# CLI helpers
-# ---------------------------
+# ---------------------------------------------------------------------------
 PRESET_QUESTIONS = [
-    "What is the deadline for the final project?",  # likely in metadata
+    "What is the deadline for the final project?",  # should hit metadata
-    "Explain the concept of tokenization in NLP.",  # content
+    "Explain the concept of polymorphism in OOP.",  # should hit FAQ
-    "How many lectures are there in the first module?",  # content
+    "How many lectures are there in the first module?",  # metadata
 ]
 
 async def run_agent(question: str, thread_id: str = "session-1"):
@@ -120,24 +139,25 @@ async def run_agent(question: str, thread_id: str = "session-1"):
         {"messages": [HumanMessage(content=question)]},
         {"configurable": {"thread_id": thread_id}},
     )
-    # The last message is the agent's reply
+    return result["messages"][-1].content
-    reply = result["messages"][-1].content
-    print(f"\nQ: {question}\nA: {reply}\n")
 
 async def main():
-    # Load data into Chroma if not already persisted
+    # Ensure FAQ is loaded once
-    if not CHROMA_DIR.exists() or not any(CHROMA_DIR.iterdir()):
+    await agent.ainvoke(
-        print("Loading FAQ data into ChromaDB...")
+        {"messages": [HumanMessage(content="load_faq_to_chroma()")]},
-        load_faq_to_chroma()
+        {"configurable": {"thread_id": "init"}},
-    print("\n--- Predefined questions ---")
+    )
+    print("\n--- Preset questions ---")
     for q in PRESET_QUESTIONS:
-        await run_agent(q)
+        ans = await run_agent(q)
+        print(f"Q: {q}\nA: {ans}\n")
     print("\n--- Interactive mode (type 'exit' to quit) ---")
     while True:
-        user_input = input("You: ")
+        user_input = input("> ")
         if user_input.lower() in {"exit", "quit"}:
             break
-        await run_agent(user_input)
+        ans = await run_agent(user_input)
+        print(ans)
 
 if __name__ == "__main__":
     asyncio.run(main())