fix: main.py — Агент с RAG-памятью

main.py

@@ -1,21 +1,20 @@
 import os
 import asyncio
-from pathlib import Path
-from typing import List
-
 from dotenv import load_dotenv
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
-from langchain_chroma import Chroma
+from langchain_core.messages import HumanMessage
 from langchain_core.documents import Document
-from langchain.text_splitter import RecursiveCharacterTextSplitter
 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 langchain_qdrant import QdrantVectorStore
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from qdrant_client import QdrantClient
 
+# Загрузка переменных окружения
 load_dotenv()
 
-# ---------- LLM ----------
+# Инициализация LLM
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -23,147 +22,90 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ---------- Vector Store ----------
+# Инициализация эмбеддингов
 embeddings = OpenAIEmbeddings(
     model="text-embedding-3-small",
     base_url="https://openrouter.ai/api/v1",
     api_key=os.getenv("OPENAI_API_KEY"),
 )
 
-vector_store = Chroma(
+# Инициализация Qdrant
-    collection_name="knowledge",
+client = QdrantClient(url="http://localhost:6333")
-    embedding_function=embeddings,
+collection_name = "knowledge_base"
+vector_store = QdrantVectorStore(
+    client=client,
+    collection_name=collection_name,
+    embeddings=embeddings,
 )
 
-# ---------- Text Splitter ----------
+# Чанкинг
-splitter = RecursiveCharacterTextSplitter(
+splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    chunk_size=1000,
-    chunk_overlap=200,
-    separators=["\n\n", "\n", " "],
-)
 
-# ---------- RAG Tools ----------
+# Инструмент поиска
 @tool
-def search_knowledge_base(query: str, max_results: int = 3) -> str:
+def search_knowledge_base(query: str, max_results: int) -> str:
-    """
+    """Semantic search in the knowledge base."""
-    Perform a semantic search in the knowledge base.
+    docs = vector_store.similarity_search(query, k=max_results)
-    Returns the concatenated contents of the most relevant documents.
-    """
-    docs: List[Document] = vector_store.similarity_search(query, k=max_results)
     if not docs:
-        return "No relevant documents found."
+        return "No results found."
-    return "\n---\n".join(doc.page_content for doc in docs)
+    return "\n".join(doc.page_content for doc in docs)
-
 
+# Инструмент добавления
 @tool
-def add_to_knowledge_base(content: str, title: str = "document") -> str:
+def add_to_knowledge_base(content: str, title: str) -> str:
-    """
+    """Add content to the knowledge base."""
-    Add a new document to the knowledge base.
-    The content will be split into chunks before indexing.
-    """
     chunks = splitter.split_text(content)
-    docs = [
+    docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks]
-        Document(page_content=chunk, metadata={"title": title, "chunk_index": i})
-        for i, chunk in enumerate(chunks)
-    ]
     vector_store.add_documents(docs)
-    return f"Added {len(docs)} chunks from '{title}' to the knowledge base."
+    return f"Added {len(docs)} chunks for {title}."
 
+# Backend для deepagents
+backend = CompositeBackend([
+    LocalShellBackend(workspace_dir="./workspace"),
+    FilesystemBackend(),
+])
 
-# ---------- Backend ----------
+# Создание агента
-backend = CompositeBackend(
-    [
-        LocalShellBackend(workspace_dir="./workspace"),
-        FilesystemBackend(),
-    ]
-)
-
-# ---------- Agent ----------
 agent = create_deep_agent(
     model=llm,
     tools=[search_knowledge_base, add_to_knowledge_base],
     backend=backend,
-    system_prompt=(
+    system_prompt="You are a helpful agent with access to a knowledge base. Use the tools to search and add information.",
-        "You are an AI assistant with access to a local knowledge base. "
-        "When you need factual information, use the provided tools: "
-        "`search_knowledge_base` to retrieve data and `add_to_knowledge_base` to store new documents. "
-        "Always cite sources from the knowledge base in your answers."
-    ),
 )
 
-# ---------- Helper Functions ----------
+async def main():
-def load_documents_from_directory(directory: Path) -> None:
+    print("Interactive agent. Commands: /add, /search, /quit")
-    """
-    Recursively read .txt files from the given directory and add them to the knowledge base.
-    """
-    for file_path in directory.rglob("*.txt"):
-        try:
-            content = file_path.read_text(encoding="utf-8")
-            title = file_path.stem
-            add_to_knowledge_base(content, title)
-            print(f"Loaded {file_path}")
-        except Exception as e:
-            print(f"Failed to load {file_path}: {e}")
-
-
-async def chat_loop() -> None:
-    """
-    Simple CLI loop.
-    Commands:
-      /add <path>   - add a text file or all txt files in a directory
-      /search <q>   - search the knowledge base
-      /quit         - exit
-    Anything else is sent to the agent as a user message.
-    """
-    thread_id = "cli-session"
-    print("AI assistant ready. Type /quit to exit.")
     while True:
-        user_input = input(">>> ").strip()
+        try:
+            user_input = input("> ").strip()
+        except EOFError:
+            break
         if not user_input:
             continue
-        if user_input.lower() == "/quit":
+        if user_input.startswith("/add"):
+            title = input("Title: ").strip()
+            content = input("Content: ").strip()
+            result = add_to_knowledge_base(content, title)
+            print(result)
+        elif user_input.startswith("/search"):
+            query = input("Query: ").strip()
+            max_str = input("Max results (int): ").strip()
+            try:
+                max_results = int(max_str)
+            except ValueError:
+                max_results = 3
+            result = search_knowledge_base(query, max_results)
+            print(result)
+        elif user_input.startswith("/quit"):
             print("Goodbye!")
             break
-        if user_input.startswith("/add"):
+        else:
-            parts = user_input.split(maxsplit=1)
+            # обычный диалог с агентом
-            if len(parts) != 2:
-                print("Usage: /add <path>")
-                continue
-            path = Path(parts[1]).expanduser().resolve()
-            if path.is_dir():
-                load_documents_from_directory(path)
-            elif path.is_file() and path.suffix.lower() == ".txt":
-                content = path.read_text(encoding="utf-8")
-                add_to_knowledge_base(content, path.stem)
-                print(f"Added file {path}")
-            else:
-                print("Provide a .txt file or a directory containing .txt files.")
-            continue
-        if user_input.startswith("/search"):
-            parts = user_input.split(maxsplit=1)
-            if len(parts) != 2:
-                print("Usage: /search <query>")
-                continue
-            query = parts[1]
-            result = search_knowledge_base(query)
-            print(f"Search results:\n{result}")
-            continue
-
-        # Normal conversation with the agent
-        try:
             response = await agent.ainvoke(
                 {"messages": [HumanMessage(content=user_input)]},
-                {"configurable": {"thread_id": thread_id}},
+                {"configurable": {"thread_id": "session-1"}},
             )
-            answer = response["messages"][-1].content
+            print(response["messages"][-1].content)
-            print(answer)
-        except Exception as e:
-            print(f"Agent error: {e}")
-
 
 if __name__ == "__main__":
-    # Optional: preload a default docs folder
+    asyncio.run(main())
-    default_dir = Path("./docs")
-    if default_dir.is_dir():
-        load_documents_from_directory(default_dir)
-    asyncio.run(chat_loop())