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

main.py

@@ -1,35 +1,29 @@
 """
-# main.py – RAG‑agent with Qdrant, OpenRouter, and LangChain
+DeepAgents RAG Agent with Qdrant and Ollama embeddings
-# -----------------------------------------------------------------
+"""
-# This script implements a simple RAG agent that can search and add
-# documents to a Qdrant vector store.  The agent is built with
-# LangChain's `create_agent` and uses OpenRouter for both the LLM and
-# embeddings.  The code follows the "Исправить" section of the
-# assignment and includes detailed comments explaining design choices.
-# -----------------------------------------------------------------
-
 import os
 import asyncio
-import argparse
 from pathlib import Path
+from typing import List
 
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
 from langchain.tools import tool
-from langchain_community.document_loaders import TextLoader
+from langchain_ollama.embeddings import OllamaEmbeddings
-from langchain_community.document_loaders import DirectoryLoader
+from langchain_qdrant import QdrantVectorStore
-from langchain_community.vectorstores import Qdrant
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.agents import create_agent, AgentExecutor, AgentType
+from deepagents import create_deep_agent
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 
-# -----------------------------------------------------------------
+# ---------------------------------------------------------------------------
-# Configuration – all secrets are read from environment variables.
+# Configuration
-# -----------------------------------------------------------------
+# ---------------------------------------------------------------------------
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+# Environment variables
-if not OPENAI_API_KEY:
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")  # required for OpenRouter LLM
-    raise RuntimeError("OPENAI_API_KEY environment variable is required")
+OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
+QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
 
-# LLM – OpenRouter gpt-oss-20b:free (free tier)
+# LLM via OpenRouter
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -37,150 +31,112 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# Embeddings – OpenAI text-embedding-3-small via OpenRouter
+# Embeddings via Ollama
-embeddings = OpenAIEmbeddings(
+embeddings = OllamaEmbeddings(model="nomic-embed-text", base_url=OLLAMA_BASE_URL)
-    model="text-embedding-3-small",
-    base_url="https://openrouter.ai/api/v1",
-    api_key=OPENAI_API_KEY,
-)
 
-# Qdrant client – assumes a local Qdrant instance running on default port
+# Qdrant vector store
-qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333")
+qdrant_vector_store = QdrantVectorStore(
-vector_store = Qdrant(
+    client_kwargs={"url": QDRANT_URL},
-    client=None,  # will be created lazily by Qdrant wrapper
+    collection_name="knowledge_base",
-    collection_name="knowledge",
     embeddings=embeddings,
-    url=qdrant_url,
+    # Create collection if not exists
+    create_collection=True,
+    # Use cosine similarity
+    distance="cosine",
 )
 
-# -----------------------------------------------------------------
+# Text splitter
-# Tool definitions – these are the only tools the agent can use.
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-# -----------------------------------------------------------------
+
+# ---------------------------------------------------------------------------
+# Tools
+# ---------------------------------------------------------------------------
 @tool
 def search_knowledge_base(query: str, max_results: int = 3) -> str:
-    """Search the knowledge base for relevant information.
+    """Semantic search in the knowledge base."""
-
+    docs = qdrant_vector_store.similarity_search(query, k=max_results)
-    Parameters
-    ----------
-    query: str
-        The search query.
-    max_results: int, optional
-        Number of top results to return (default 3).
-    """
-    docs = vector_store.similarity_search(query, k=max_results)
     if not docs:
         return "No results found."
-    return "\n\n---\n\n".join([f"{doc.metadata.get('title', 'Untitled')}\n{doc.page_content}" for doc in docs])
+    return "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs))
 
 @tool
-def add_to_knowledge_base(content: str, title: str = "Untitled") -> str:
+def add_to_knowledge_base(content: str, title: str = "document") -> str:
-    """Add a new document (or chunk) to the knowledge base.
+    """Add content to the knowledge base after chunking."""
+    # Split content into chunks
+    chunks = text_splitter.split_text(content)
+    metadatas = [{"title": title, "chunk_index": i} for i in range(len(chunks))]
+    # Add to Qdrant
+    qdrant_vector_store.add_texts(chunks, metadatas=metadatas)
+    return f"Added {len(chunks)} chunks from '{title}'."
 
-    Parameters
+# ---------------------------------------------------------------------------
-    ----------
+# Backend setup
-    content: str
+# ---------------------------------------------------------------------------
-        The text content to add.
+backend = CompositeBackend([
-    title: str, optional
+    LocalShellBackend(workspace_dir="./workspace"),
-        A human‑readable title for the document.
+    FilesystemBackend(),
-    """
+])
-    doc = {
-        "page_content": content,
-        "metadata": {"title": title},
-    }
-    vector_store.add_documents([doc])
-    return f"Added document '{title}'."
 
-# -----------------------------------------------------------------
+# ---------------------------------------------------------------------------
-# Agent setup – using LangChain's create_agent with a custom system prompt.
+# Agent creation
-# -----------------------------------------------------------------
+# ---------------------------------------------------------------------------
-SYSTEM_PROMPT = (
+agent = create_deep_agent(
-    "You are a helpful assistant with access to a knowledge base. "
+    model=llm,
-    "Use the tools `search_knowledge_base` and `add_to_knowledge_base` "
-    "to answer user queries.  If the user asks to add information, "
-    "use `add_to_knowledge_base`.  If the user asks for information, "
-    "use `search_knowledge_base`.  Do not fabricate facts."
-)
-
-agent = create_agent(
-    llm=llm,
     tools=[search_knowledge_base, add_to_knowledge_base],
-    system_prompt=SYSTEM_PROMPT,
+    backend=backend,
-    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+    system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.",
 )
 
-executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
+# ---------------------------------------------------------------------------
+# Document loader for initialization
+# ---------------------------------------------------------------------------
+async def load_documents_from_dir(directory: str):
+    """Load all .txt files from a directory into the knowledge base."""
+    dir_path = Path(directory)
+    for file_path in dir_path.rglob("*.txt"):
+        content = file_path.read_text(encoding="utf-8")
+        title = file_path.stem
+        await agent.ainvoke(
+            {"messages": [HumanMessage(content=f"/add {content}")], "metadata": {"title": title}},
+            {"configurable": {"thread_id": "init-session"}},
+        )
 
-# -----------------------------------------------------------------
+# ---------------------------------------------------------------------------
-# Document ingestion – split into chunks and add to Qdrant.
+# CLI client
-# -----------------------------------------------------------------
+# ---------------------------------------------------------------------------
-def ingest_directory(directory: str, chunk_size: int = 1000, chunk_overlap: int = 200):
+async def cli():
-    """Load all text files from *directory*, split into chunks, and store.
+    print("DeepAgents RAG CLI. Commands: /add <text>, /search <query>, /quit")
-
-    Parameters
-    ----------
-    directory: str
-        Path to the directory containing documents.
-    chunk_size: int, optional
-        Size of each chunk in characters.
-    chunk_overlap: int, optional
-        Overlap between consecutive chunks.
-    """
-    loader = DirectoryLoader(directory, glob="**/*.txt")
-    documents = loader.load()
-    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-    chunks = splitter.split_documents(documents)
-    # Convert LangChain Document objects to dicts expected by Qdrant
-    docs_to_add = []
-    for doc in chunks:
-        title = doc.metadata.get("source", "Untitled")
-        docs_to_add.append({
-            "page_content": doc.page_content,
-            "metadata": {"title": title, "source": doc.metadata.get("source", "")},
-        })
-    vector_store.add_documents(docs_to_add)
-    print(f"Ingested {len(docs_to_add)} chunks into the knowledge base.")
-
-# -----------------------------------------------------------------
-# CLI – simple interactive loop.
-# -----------------------------------------------------------------
-async def main():
-    parser = argparse.ArgumentParser(description="RAG Agent CLI")
-    parser.add_argument("--ingest", type=str, help="Path to directory to ingest")
-    args = parser.parse_args()
-
-    if args.ingest:
-        ingest_directory(args.ingest)
-        return
-
-    print("RAG Agent ready. Type /quit to exit.")
     while True:
-        user_input = input("You: ")
+        user_input = input("> ")
         if user_input.strip() == "/quit":
             print("Goodbye!")
             break
         if user_input.startswith("/add "):
-            # Expected format: /add <title> | <content>
+            content = user_input[5:].strip()
-            try:
+            response = await agent.ainvoke(
-                _, rest = user_input.split("/add ", 1)
+                {"messages": [HumanMessage(content=f"/add {content}")], "metadata": {"title": "user_input"}},
-                title, content = rest.split("|", 1)
+                {"configurable": {"thread_id": "cli-session"}},
-                title = title.strip()
+            )
-                content = content.strip()
+            print(response["messages"][-1].content)
-                result = await executor.ainvoke({"messages": [HumanMessage(content=f"Add document {title}")], "configurable": {"thread_id": "session-1"}})
+        elif user_input.startswith("/search "):
-                # Directly call tool to add content
+            query = user_input[8:].strip()
-                add_to_knowledge_base(content, title)
+            response = await agent.ainvoke(
-                print("Agent: Document added.")
+                {"messages": [HumanMessage(content=f"/search {query}")], "metadata": {"query": query}},
-            except Exception as e:
+                {"configurable": {"thread_id": "cli-session"}},
-                print(f"Error parsing /add command: {e}")
+            )
-            continue
+            print(response["messages"][-1].content)
-        if user_input.startswith("/search "):
+        else:
-            query = user_input[len("/search "):].strip()
+            print("Unknown command. Use /add, /search, or /quit.")
-            result = await executor.ainvoke({"messages": [HumanMessage(content=f"Search for {query}")], "configurable": {"thread_id": "session-1"}})
+
-            print("Agent:", result["messages"][-1].content)
+# ---------------------------------------------------------------------------
-            continue
+# Main entry point
-        # Default: normal chat
+# ---------------------------------------------------------------------------
-        result = await executor.ainvoke({"messages": [HumanMessage(content=user_input)], "configurable": {"thread_id": "session-1"}})
+async def main():
-        print("Agent:", result["messages"][-1].content)
+    # Optional: load initial documents from a folder named 'data'
+    data_dir = Path("./data")
+    if data_dir.exists():
+        await load_documents_from_dir(str(data_dir))
+    await cli()
 
 if __name__ == "__main__":
     asyncio.run(main())
-"""