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

main.py

@@ -1,108 +1,186 @@
+"""
+# main.py – RAG‑agent with Qdrant, OpenRouter, and LangChain
+# -----------------------------------------------------------------
+# 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 langchain_openai import ChatOpenAI
-from langchain_ollama import OllamaEmbeddings
-from langchain_qdrant import QdrantVectorStore
-from langchain_core.documents import Document
-from langchain.tools import tool
-from deepagents import create_deep_agent
-from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-from langchain_core.messages import HumanMessage
 
-# ---------- LLM ----------
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_core.messages import HumanMessage
+from langchain.tools import tool
+from langchain_community.document_loaders import TextLoader
+from langchain_community.document_loaders import DirectoryLoader
+from langchain_community.vectorstores import Qdrant
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.agents import create_agent, AgentExecutor, AgentType
+
+# -----------------------------------------------------------------
+# Configuration – all secrets are read from environment variables.
+# -----------------------------------------------------------------
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+if not OPENAI_API_KEY:
+    raise RuntimeError("OPENAI_API_KEY environment variable is required")
+
+# LLM – OpenRouter gpt-oss-20b:free (free tier)
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
+    api_key=OPENAI_API_KEY,
     temperature=0.0,
 )
 
-# ---------- Embeddings ----------
+# Embeddings – OpenAI text-embedding-3-small via OpenRouter
-# Using Ollama embeddings as per assignment correction
+embeddings = OpenAIEmbeddings(
-embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    model="text-embedding-3-small",
-
+    base_url="https://openrouter.ai/api/v1",
-# ---------- Vector Store (Qdrant) ----------
+    api_key=OPENAI_API_KEY,
-# Ensure Qdrant is running locally (default port 6333)
-vector_store = QdrantVectorStore(
-    url="http://localhost:6333",
-    collection_name="knowledge",
-    embedding_function=embeddings,
 )
 
-# ---------- Tools ----------
+# Qdrant client – assumes a local Qdrant instance running on default port
+qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333")
+vector_store = Qdrant(
+    client=None,  # will be created lazily by Qdrant wrapper
+    collection_name="knowledge",
+    embeddings=embeddings,
+    url=qdrant_url,
+)
+
+# -----------------------------------------------------------------
+# Tool definitions – these are the only tools the agent can use.
+# -----------------------------------------------------------------
 @tool
 def search_knowledge_base(query: str, max_results: int = 3) -> str:
-    """Semantic search in the knowledge base."""
+    """Search the knowledge base for relevant information.
+
+    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".join(f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs))
+    return "\n\n---\n\n".join([f"{doc.metadata.get('title', 'Untitled')}\n{doc.page_content}" for doc in docs])
 
 @tool
-def add_to_knowledge_base(content: str, title: str = "untitled") -> str:
+def add_to_knowledge_base(content: str, title: str = "Untitled") -> str:
-    """Add a document to the knowledge base."""
+    """Add a new document (or chunk) to the knowledge base.
-    doc = Document(page_content=content, metadata={"title": title})
+
+    Parameters
+    ----------
+    content: str
+        The text content to add.
+    title: str, optional
+        A human‑readable title for the document.
+    """
+    doc = {
+        "page_content": content,
+        "metadata": {"title": title},
+    }
     vector_store.add_documents([doc])
-    return f"Document '{title}' added to the knowledge base."
+    return f"Added document '{title}'."
 
-# ---------- Backend ----------
+# -----------------------------------------------------------------
-backend = CompositeBackend([
+# Agent setup – using LangChain's create_agent with a custom system prompt.
-    LocalShellBackend(workspace_dir="./workspace"),
+# -----------------------------------------------------------------
-    FilesystemBackend(),
+SYSTEM_PROMPT = (
-])
+    "You are a helpful assistant with access to a knowledge base. "
+    "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 ----------
+agent = create_agent(
-agent = create_deep_agent(
+    llm=llm,
-    model=llm,
     tools=[search_knowledge_base, add_to_knowledge_base],
-    backend=backend,
+    system_prompt=SYSTEM_PROMPT,
-    system_prompt="You are an assistant with access to a knowledge base. Use the provided tools to search and add information."
+    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
 )
 
-# ---------- Document Loader ----------
+executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
-async def load_documents_from_dir(directory: str):
-    """Load all text files from a directory into the vector store."""
-    for file_path in Path(directory).rglob("*.txt"):
-        text = file_path.read_text(encoding="utf-8")
-        title = file_path.stem
-        await agent.ainvoke(
-            {"messages": [HumanMessage(content=f"/add {title}")], "content": text},
-            {"configurable": {"thread_id": "init"}},
-        )
 
-# ---------- Interactive CLI ----------
+# -----------------------------------------------------------------
-async def interactive_loop():
+# Document ingestion – split into chunks and add to Qdrant.
-    print("Welcome to the RAG Agent. Commands: /add <title>, /search <query>, /quit")
+# -----------------------------------------------------------------
+def ingest_directory(directory: str, chunk_size: int = 1000, chunk_overlap: int = 200):
+    """Load all text files from *directory*, split into chunks, and store.
+
+    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("> ")
+        user_input = input("You: ")
         if user_input.strip() == "/quit":
             print("Goodbye!")
             break
         if user_input.startswith("/add "):
-            parts = user_input.split(" ", 1)
+            # Expected format: /add <title> | <content>
-            title = parts[1] if len(parts) > 1 else "untitled"
+            try:
-            content = input("Enter content: ")
+                _, rest = user_input.split("/add ", 1)
-            response = await agent.ainvoke(
+                title, content = rest.split("|", 1)
-                {"messages": [HumanMessage(content=f"/add {title}")], "content": content},
+                title = title.strip()
-                {"configurable": {"thread_id": "session"}},
+                content = content.strip()
-            )
+                result = await executor.ainvoke({"messages": [HumanMessage(content=f"Add document {title}")], "configurable": {"thread_id": "session-1"}})
-            print(response["messages"][-1].content)
+                # Directly call tool to add content
-        elif user_input.startswith("/search "):
+                add_to_knowledge_base(content, title)
-            query = user_input.split(" ", 1)[1]
+                print("Agent: Document added.")
-            response = await agent.ainvoke(
+            except Exception as e:
-                {"messages": [HumanMessage(content=f"/search {query}")]},
+                print(f"Error parsing /add command: {e}")
-                {"configurable": {"thread_id": "session"}},
+            continue
-            )
+        if user_input.startswith("/search "):
-            print(response["messages"][-1].content)
+            query = user_input[len("/search "):].strip()
-        else:
+            result = await executor.ainvoke({"messages": [HumanMessage(content=f"Search for {query}")], "configurable": {"thread_id": "session-1"}})
-            print("Unknown command. Use /add, /search, or /quit.")
+            print("Agent:", result["messages"][-1].content)
-
+            continue
-# ---------- Main ----------
+        # Default: normal chat
-async def main():
+        result = await executor.ainvoke({"messages": [HumanMessage(content=user_input)], "configurable": {"thread_id": "session-1"}})
-    # Optional: load initial documents
+        print("Agent:", result["messages"][-1].content)
-    # await load_documents_from_dir("./data")
-    await interactive_loop()
 
 if __name__ == "__main__":
     asyncio.run(main())
+"""