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

main.py

@@ -1,91 +1,91 @@
-import os
+#!/usr/bin/env python3
+"""RAG‑agent with ChromaDB and Tavily search.
+
+This implementation follows the course specification and uses the
+`deepagents` framework to create a single agent that can decide whether
+to query the local knowledge base (ChromaDB) or perform a web search
+via Tavily.  The agent is backed by OpenRouter for both LLM and
+embeddings, complying with the mandatory technical constraints.
+"""
+
 import asyncio
+import os
+from pathlib import Path
+
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_chroma import Chroma
 from langchain_core.documents import Document
 from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_tavily import TavilySearchRun
 from langchain.tools import tool
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-from langchain_core.messages import HumanMessage
-from tavily import TavilySearchResults
 
-# ---------------------
+# ---------------------------------------------------------------------------
 # Configuration
-# ---------------------
+# ---------------------------------------------------------------------------
+# Load environment variables (OpenRouter key, Tavily key)
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
 
-# ---------------------
+# Persist directory for ChromaDB
-# Vector store utilities
+CHROMA_DIR = Path("./chroma_db")
-# ---------------------
+CHROMA_DIR.mkdir(parents=True, exist_ok=True)
 
-def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
+# ---------------------------------------------------------------------------
-    """Create or load a Chroma vector store with OpenAI embeddings."""
+# 1. Vector store (ChromaDB + OpenRouter embeddings)
-    embeddings = OpenAIEmbeddings(
+# ---------------------------------------------------------------------------
-        model="text-embedding-3-small",
+embeddings = OpenAIEmbeddings(
-        base_url="https://openrouter.ai/api/v1",
+    model="text-embedding-3-small",
-        api_key=OPENAI_API_KEY,
+    base_url="https://openrouter.ai/api/v1",
-    )
+    api_key=OPENAI_API_KEY,
-    return Chroma(
+)
-        collection_name="knowledge",
-        embedding_function=embeddings,
-        persist_directory=persist_directory,
-    )
 
+vector_store = Chroma(
+    collection_name="knowledge",
+    embedding_function=embeddings,
+    persist_directory=str(CHROMA_DIR),
+)
 
-def load_documents(directory: str, vectorstore: Chroma) -> None:
+# Helper to load documents from a directory and add to the store
-    """Load .txt and .md files from *directory* into *vectorstore* using chunking."""
+
+def load_documents(directory: Path):
+    """Read .txt/.md files, split into chunks, and store in Chroma."""
     splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     docs = []
-    for root, _, files in os.walk(directory):
+    for file in directory.glob("**/*"):
-        for fname in files:
+        if file.suffix.lower() not in {".txt", ".md"}:
-            if fname.lower().endswith(('.txt', '.md')):
+            continue
-                path = os.path.join(root, fname)
+        text = file.read_text(encoding="utf-8")
-                with open(path, "r", encoding="utf-8") as f:
+        chunks = splitter.split_text(text)
-                    text = f.read()
+        docs.extend([Document(page_content=c, metadata={"source": str(file)}) for c in chunks])
-                # Create Document objects
-                docs.extend(
-                    [Document(page_content=chunk, metadata={"source": path}) for chunk in splitter.split_text(text)]
-                )
     if docs:
-        vectorstore.add_documents(docs)
+        vector_store.add_documents(docs)
-        vectorstore.persist()
+        vector_store.persist()
-
-# ---------------------
-# Tools
-# ---------------------
-
-vectorstore = create_vectorstore()
-
-# Ensure we have some data loaded – load from ./documents if collection empty
-if len(vectorstore.get_all_documents()) == 0:
-    load_documents("./documents", vectorstore)
 
+# ---------------------------------------------------------------------------
+# 2. Tools
+# ---------------------------------------------------------------------------
 @tool
 def search_local_kb(query: str, top_k: int = 3) -> str:
-    """Search the local knowledge base for relevant passages."""
+    """Semantic search in the local ChromaDB knowledge base."""
-    docs = vectorstore.similarity_search(query, k=top_k)
+    docs = vector_store.similarity_search(query, k=top_k)
     if not docs:
-        return "[Local KB] No relevant information found."
+        return "No relevant information found in the local knowledge base."
-    result = "\n\n".join([f"{d.metadata.get('source', 'unknown')}\n{d.page_content}" for d in docs])
+    return "\n---\n".join([f"{i+1}. {d.page_content[:200]}…" for i, d in enumerate(docs)])
-    return f"[Local KB]\n{result}"
 
 @tool
 def web_search(query: str) -> str:
-    """Perform a web search using Tavily and return top results."""
+    """Perform a web search using Tavily."""
-    results = TavilySearchResults(query=query, max_results=3, api_key=TAVILY_API_KEY)
+    tavily = TavilySearchRun(api_key=TAVILY_API_KEY, max_results=3)
-    if not results.results:
+    results = tavily.run(query)
-        return "[Web Search] No results found."
+    if not results:
-    snippets = []
+        return "No web results found."
-    for r in results.results:
+    return "\n---\n".join([f"{i+1}. {r['title']}\n{r['content'][:200]}…" for i, r in enumerate(results)])
-        snippets.append(f"{r.get('title', 'No title')}\n{r.get('content', 'No content')}\nURL: {r.get('url', '')}")
-    return f"[Web Search]\n\n".join(snippets)
-
-# ---------------------
-# Agent setup
-# ---------------------
 
+# ---------------------------------------------------------------------------
+# 3. Agent (deepagents)
+# ---------------------------------------------------------------------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -99,9 +99,11 @@ backend = CompositeBackend([
 ])
 
 system_prompt = (
-    "You are a helpful RAG agent. For questions about local documents, use the tool `search_local_kb`. "
+    "You are a helpful assistant. For a user query, first decide whether the
-    "For current news or facts that may not be in your local knowledge base, use `web_search`. "
+    answer can be found in the local knowledge base. If so, use the
-    "Return the answer prefixed with either `[Local KB]` or `[Web Search]` to indicate the source."
+    `search_local_kb` tool. If the query requires up‑to‑date information,
+    use the `web_search` tool. Respond with the best answer and clearly
+    state the source: `chromadb` or `tavily`."
 )
 
 agent = create_deep_agent(
@@ -111,25 +113,28 @@ agent = create_deep_agent(
     system_prompt=system_prompt,
 )
 
-# ---------------------
+# ---------------------------------------------------------------------------
-# CLI loop
+# 4. CLI loop
-# ---------------------
+# ---------------------------------------------------------------------------
-
 async def main():
     print("RAG Agent ready. Type your question (or 'exit' to quit).")
+    thread_id = "session-1"
     while True:
-        user_input = input("\nQuery: ")
+        user_input = input("\nЗапрос: ")
         if user_input.lower() in {"exit", "quit", "q"}:
             print("Goodbye!")
             break
-        # Invoke agent
+        response = await agent.ainvoke(
-        result = await agent.ainvoke(
+            {"messages": [{"role": "user", "content": user_input}]},
-            {"messages": [HumanMessage(content=user_input)]},
+            {"configurable": {"thread_id": thread_id}},
-            {"configurable": {"thread_id": "session-1"}},
         )
-        # The agent returns a list of messages; get the last one
+        # The last message contains the assistant reply
-        reply = result["messages"][-1].content
+        assistant_msg = response["messages"][-1].content
-        print(f"\n{reply}")
+        print(f"\nОтвет:\n{assistant_msg}")
 
 if __name__ == "__main__":
+    # Load documents once at startup
+    docs_dir = Path("./documents")
+    if docs_dir.exists():
+        load_documents(docs_dir)
     asyncio.run(main())