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

main.py

@@ -1,89 +1,26 @@
-import os
 import asyncio
+import os
+import json
 from pathlib import Path
+
 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
 
-# ===================== CONFIG =====================
+# ------------------------------------------------------------------
+# Configuration
+# ------------------------------------------------------------------
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 if not OPENAI_API_KEY:
-    raise RuntimeError("OPENAI_API_KEY is not set in environment")
+    raise RuntimeError("OPENAI_API_KEY not set in environment")
 
-# ===================== EMBEDDINGS & VECTOR STORE =====================
+# ------------------------------------------------------------------
-embeddings = OpenAIEmbeddings(
+# LLM and embeddings (OpenRouter only)
-    model="text-embedding-3-small",
+# ------------------------------------------------------------------
-    base_url="https://openrouter.ai/api/v1",
-    api_key=OPENAI_API_KEY,
-)
-
-vector_store = Chroma(
-    collection_name="faq_knowledge",
-    embedding_function=embeddings,
-    persist_directory="./chroma_faq",
-)
-
-# ===================== TOOL: SEARCH IN CHROMA =====================
-@tool
-def search_course_docs(query: str, k: int = 3) -> str:
-    """Search the local FAQ knowledge base for relevant passages."""
-    docs: list[Document] = vector_store.similarity_search(query, k=k)
-    if not docs:
-        return "No relevant information found in the course materials."
-    return "\n\n---\n\n".join(f"**{d.metadata.get('title', 'Untitled')}**\n{d.page_content}" for d in docs)
-
-# ===================== TOOL: FETCH METADATA (MCP‑STYLE) =====================
-# For the purpose of this assignment we use a static JSON file as the mock MCP server response.
-METADATA_JSON = {
-    "schedule": {
-        "Monday": "Lecture 1: Introduction",
-        "Wednesday": "Lecture 2: Advanced Topics",
-        "Friday": "Lecture 3: Practical Applications"
-    },
-    "instructors": {
-        "Dr. Smith": "smith@example.com",
-        "Prof. Doe": "doe@example.com"
-    }
-}
-
-@tool
-def fetch_course_meta(query: str) -> str:
-    """Mock MCP tool that returns course metadata based on the query.
-    In production this would be an HTTP GET to an MCP server.
-    """
-    query = query.lower()
-    if "schedule" in query:
-        return "\n".join(f"{day}: {info}" for day, info in METADATA_JSON["schedule"].items())
-    if "instructor" in query or "email" in query:
-        return "\n".join(f"{name}: {email}" for name, email in METADATA_JSON["instructors"].items())
-    return "No metadata matches your query."
-
-# ===================== LOAD FAQ TO CHROMA =====================
-MD_DIR = Path("data")
-if not MD_DIR.exists():
-    MD_DIR.mkdir(parents=True, exist_ok=True)
-    # Create example markdown files if none exist
-    (MD_DIR / "faq1.md").write_text("# FAQ 1\nWhat is the course about?\nThe course covers advanced AI techniques.")
-    (MD_DIR / "faq2.md").write_text("# FAQ 2\nHow to install dependencies?\nRun `pip install -r requirements.txt`.")
-    (MD_DIR / "faq3.md").write_text("# FAQ 3\nWhere to find the schedule?\nCheck the course website.")
-
-# Chunking and persisting
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-
-for md_file in MD_DIR.glob("*.md"):
-    content = md_file.read_text(encoding="utf-8")
-    docs = [Document(page_content=chunk, metadata={"title": md_file.stem}) for chunk in text_splitter.split_text(content)]
-    vector_store.add_documents(docs)
-
-vector_store.persist()
-
-# ===================== AGENT SETUP =====================
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -91,44 +28,135 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
+embeddings = OpenAIEmbeddings(
+    model="text-embedding-3-small",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=OPENAI_API_KEY,
+)
+
+# ------------------------------------------------------------------
+# Chroma vector store (persisted)
+# ------------------------------------------------------------------
+CHROMA_PATH = Path("./chroma_faq")
+vector_store = Chroma(
+    collection_name="faq_collection",
+    embedding_function=embeddings,
+    persist_directory=str(CHROMA_PATH),
+)
+
+# ------------------------------------------------------------------
+# Utility: load markdown files into Chroma
+# ------------------------------------------------------------------
+async def load_faq_to_chroma(md_dir: str = "data"):
+    md_dir = Path(md_dir)
+    if not md_dir.is_dir():
+        raise FileNotFoundError(f"Markdown directory {md_dir} not found")
+    docs = []
+    for md_file in md_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()
+    print(f"Loaded {len(docs)} documents into Chroma (persisted at {CHROMA_PATH})")
+
+# ------------------------------------------------------------------
+# Tools
+# ------------------------------------------------------------------
+@tool
+def search_course_docs(query: str) -> str:
+    """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."
+    return "\n---\n".join(f"[{doc.metadata.get('source', 'unknown')}] {doc.page_content}" for doc in results)
+
+@tool
+def fetch_course_meta(query: str) -> str:
+    """Simulate an MCP-style tool that fetches course metadata.
+    In production this would be a real HTTP call to an MCP server.
+    Here we use a local JSON file as a mock response.
+    """
+    meta_file = Path("meta.json")
+    if not meta_file.is_file():
+        return "Metadata source not available."
+    data = json.loads(meta_file.read_text(encoding="utf-8"))
+    # Very naive search: return items where query string appears in any value
+    matches = []
+    for key, value in data.items():
+        if isinstance(value, str) and query.lower() in value.lower():
+            matches.append(f"{key}: {value}")
+    if not matches:
+        return "No metadata matches found."
+    return "\n".join(matches)
+
+# ------------------------------------------------------------------
+# Backend setup
+# ------------------------------------------------------------------
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
+# ------------------------------------------------------------------
+# DeepAgent creation
+# ------------------------------------------------------------------
 agent = create_deep_agent(
     model=llm,
     tools=[search_course_docs, fetch_course_meta],
     backend=backend,
-    system_prompt="You are a helpful FAQ assistant. Use only the provided tools. In your final answer, prefix the source with `source: chroma` or `source: mcp_meta` accordingly.",
+    system_prompt=(
+        "You are a helpful FAQ assistant for the course. "
+        "When answering a question, use the local Chroma database if the answer is about course content. "
+        "If the question is about schedule, metadata, or other non-content info, call fetch_course_meta. "
+        "Always indicate the source in your final answer as either 'source: chroma' or 'source: mcp_meta'."
+    ),
 )
 
-# ===================== CLI =====================
+# ------------------------------------------------------------------
+# CLI helpers
+# ------------------------------------------------------------------
 PRESET_QUESTIONS = [
-    "What is the course about?",  # chroma
+    "What topics are covered in the first lecture?",
-    "How to install dependencies?",  # chroma
+    "Explain the concept of recursion as described in the notes.",
-    "What is the lecture schedule?",  # mcp_meta
+    "When is the next lab session scheduled?",
 ]
 
-async def run_cli():
+async def run_interactive():
-    print("=== FAQ Assistant ===")
+    print("--- FAQ Bot CLI ---")
-    print("Type your question or 'exit' to quit.")
+    print("Type 'exit' to quit.")
-    for i, q in enumerate(PRESET_QUESTIONS, 1):
-        print(f"\nPreset {i}: {q}")
-        await handle_question(q)
     while True:
-        user_input = input("\nYour question: ")
+        user_input = input("\nQuestion: ")
         if user_input.lower() in {"exit", "quit"}:
             break
-        await handle_question(user_input)
+        result = await agent.ainvoke(
+            {"messages": [HumanMessage(content=user_input)]},
+            {"configurable": {"thread_id": "session-1"}},
+        )
+        # The last message is the assistant's reply
+        reply = result["messages"][-1].content
+        print("\nAnswer:\n", reply)
 
-async def handle_question(question: str):
+async def run_presets():
-    result = await agent.ainvoke(
+    for q in PRESET_QUESTIONS:
-        {"messages": [HumanMessage(content=question)]},
+        print("\nQuestion:", q)
-        {"configurable": {"thread_id": "session-1"}},
+        result = await agent.ainvoke(
-    )
+            {"messages": [HumanMessage(content=q)]},
-    answer = result["messages"][-1].content
+            {"configurable": {"thread_id": "session-1"}},
-    print("\nAnswer:\n", answer)
+        )
+        reply = result["messages"][-1].content
+        print("Answer:\n", reply)
+
+# ------------------------------------------------------------------
+# Main entry point
+# ------------------------------------------------------------------
+async def main():
+    # Load data into Chroma if not already persisted
+    if not CHROMA_PATH.is_dir() or not any(CHROMA_PATH.iterdir()):
+        await load_faq_to_chroma()
+    # Run preset questions first
+    await run_presets()
+    # Then interactive mode
+    await run_interactive()
 
 if __name__ == "__main__":
-    asyncio.run(run_cli())
+    asyncio.run(main())