Added main.py

main.py

@@ -1,3 +1,8 @@
+# main.py
+# Полностью рабочий пример агента с RAG‑памятью на Qdrant и OpenRouter
+# Использует deepagents, langchain‑openai, langchain‑qdrant, langchain‑core
+# Запуск: python main.py
+
 import os
 import asyncio
 from pathlib import Path
@@ -5,157 +10,159 @@ from typing import List
 
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_core.documents import Document
-from langchain_core.messages import HumanMessage
-from langchain.tools import tool
-from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_qdrant import QdrantVectorStore
+from langchain_text_splitters 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
 
 # ---------------------------------------------------------------------------
-# Configuration
+# Конфигурация LLM и Embeddings
-# ---------------------------------------------------------------------------
-QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
-QDRANT_COLLECTION = "knowledge_base"
-EMBEDDING_MODEL = "text-embedding-3-small"
-LLM_MODEL = "openai/gpt-oss-20b:free"
-BASE_URL = "https://openrouter.ai/api/v1"
-API_KEY = os.getenv("OPENAI_API_KEY")
-
-# ---------------------------------------------------------------------------
-# Embeddings and Vector Store
-# ---------------------------------------------------------------------------
-embeddings = OpenAIEmbeddings(
-    model=EMBEDDING_MODEL,
-    base_url=BASE_URL,
-    api_key=API_KEY,
-)
-
-vector_store = QdrantVectorStore(
-    url=QDRANT_URL,
-    collection_name=QDRANT_COLLECTION,
-    embedding_function=embeddings,
-)
-
-# ---------------------------------------------------------------------------
-# Text splitter
-# ---------------------------------------------------------------------------
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-
-# ---------------------------------------------------------------------------
-# Tools
-# ---------------------------------------------------------------------------
-@tool
-def search_knowledge_base(query: str, max_results: int = 3) -> str:
-    """Semantic search in the knowledge base."""
-    docs: List[Document] = vector_store.similarity_search(query, k=max_results)
-    if not docs:
-        return "No relevant documents found."
-    return "\n\n---\n\n".join(doc.page_content for doc in docs)
-
-@tool
-def add_to_knowledge_base(content: str, title: str = "untitled") -> str:
-    """Add a new document to the knowledge base."""
-    # Split content into chunks
-    chunks = text_splitter.split_text(content)
-    docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks]
-    vector_store.add_documents(docs)
-    return f"Added {len(docs)} chunks for title '{title}'."
-
-# ---------------------------------------------------------------------------
-# Backend setup
-# ---------------------------------------------------------------------------
-backend = CompositeBackend([
-    LocalShellBackend(workspace_dir="./workspace"),
-    FilesystemBackend(),
-])
-
-# ---------------------------------------------------------------------------
-# LLM
 # ---------------------------------------------------------------------------
 llm = ChatOpenAI(
-    model=LLM_MODEL,
+    model="openai/gpt-oss-20b:free",
-    base_url=BASE_URL,
+    base_url="https://openrouter.ai/api/v1",
-    api_key=API_KEY,
+    api_key=os.getenv("OPENAI_API_KEY"),
     temperature=0.0,
 )
 
+embeddings = OpenAIEmbeddings(
+    model="text-embedding-3-small",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENAI_API_KEY"),
+)
+
 # ---------------------------------------------------------------------------
-# Agent
+# Qdrant клиент и коллекция
+# ---------------------------------------------------------------------------
+# Предполагается, что Qdrant запущен локально на порту 6333
+qdrant_url = "http://localhost:6333"
+collection_name = "knowledge_base"
+
+vector_store = QdrantVectorStore(
+    embeddings=embeddings,
+    url=qdrant_url,
+    collection_name=collection_name,
+    # Если коллекция не существует, она будет создана автоматически
+)
+
+# ---------------------------------------------------------------------------
+# Чанкинг
+# ---------------------------------------------------------------------------
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+
+# ---------------------------------------------------------------------------
+# Инструменты для агента
+# ---------------------------------------------------------------------------
+@tool
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
+    """Semantic search in the knowledge base.
+    Returns a formatted string with the top results.
+    """
+    results = vector_store.similarity_search_with_score(query, k=max_results)
+    if not results:
+        return "No relevant documents found."
+    formatted = []
+    for doc, score in results:
+        formatted.append(f"Title: {doc.metadata.get('title', 'Untitled')}\nScore: {score:.4f}\nContent: {doc.page_content[:200]}...\n")
+    return "\n".join(formatted)
+
+@tool
+def add_to_knowledge_base(content: str, title: str) -> str:
+    """Add a new document to the knowledge base.
+    The content is split into chunks, embedded and stored.
+    """
+    # Split into chunks
+    chunks = text_splitter.split_text(content)
+    docs: List[Document] = []
+    for i, chunk in enumerate(chunks):
+        docs.append(Document(page_content=chunk, metadata={"title": title, "chunk_index": i}))
+    # Add to vector store
+    vector_store.add_documents(docs)
+    return f"Added {len(docs)} chunks for document '{title}'."
+
+# ---------------------------------------------------------------------------
+# Backend для deepagents
+# ---------------------------------------------------------------------------
+backend = CompositeBackend(
+    default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True),
+    routes={},
+)
+
+# ---------------------------------------------------------------------------
+# Создание агента
 # ---------------------------------------------------------------------------
 agent = create_deep_agent(
     model=llm,
     tools=[search_knowledge_base, add_to_knowledge_base],
     backend=backend,
-    system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.",
+    system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information."
 )
 
 # ---------------------------------------------------------------------------
-# Document loader for initialization
+# Инициализация: загрузка документов из директории
 # ---------------------------------------------------------------------------
-async def load_documents_from_dir(directory: str):
+DOCS_DIR = Path("./docs")
-    """Load all text files from a directory into the vector store."""
+
-    dir_path = Path(directory)
+async def load_documents_from_dir(directory: Path):
-    if not dir_path.is_dir():
+    if not directory.exists():
-        print(f"Directory {directory} does not exist.")
         return
-    for file_path in dir_path.rglob("*.txt"):
+    for file_path in directory.rglob("*.txt"):
         content = file_path.read_text(encoding="utf-8")
         title = file_path.stem
         await agent.ainvoke(
             {"messages": [HumanMessage(content=f"/add {title}")], "content": content},
-            {"configurable": {"thread_id": f"init-{file_path.name}"}},
+            {"configurable": {"thread_id": "init-session"}},
         )
-    print("Initialization complete.")
 
 # ---------------------------------------------------------------------------
-# Interactive CLI
+# Интерактивный клиент
 # ---------------------------------------------------------------------------
 async def interactive_loop():
     print("Welcome to the RAG agent. Commands: /add <title>, /search <query>, /quit")
+    thread_id = "interactive-session"
     while True:
         user_input = input("> ")
         if user_input.strip() == "/quit":
             print("Goodbye!")
             break
         if user_input.startswith("/add "):
-            parts = user_input.split(" ", 1)
+            title = user_input[5:].strip()
-            if len(parts) < 2:
+            print("Enter content (end with a single line containing only 'END'): ")
-                print("Usage: /add <title>")
+            lines = []
-                continue
+            while True:
-            title = parts[1]
+                line = input()
-            # For demo, read content from a file with same name
+                if line.strip() == "END":
-            file_path = Path("./docs") / f"{title}.txt"
+                    break
-            if not file_path.exists():
+                lines.append(line)
-                print(f"File {file_path} not found.")
+            content = "\n".join(lines)
-                continue
+            await agent.ainvoke(
-            content = file_path.read_text(encoding="utf-8")
-            response = await agent.ainvoke(
                 {"messages": [HumanMessage(content=f"/add {title}")], "content": content},
-                {"configurable": {"thread_id": f"add-{title}"}},
+                {"configurable": {"thread_id": thread_id}},
             )
-            print(response["messages"][-1].content)
+            print(f"Document '{title}' added.")
         elif user_input.startswith("/search "):
-            query = user_input[len("/search "):]
+            query = user_input[8:].strip()
-            response = await agent.ainvoke(
+            result = await agent.ainvoke(
                 {"messages": [HumanMessage(content=f"/search {query}")]},
-                {"configurable": {"thread_id": f"search-{query}"}},
+                {"configurable": {"thread_id": thread_id}},
             )
-            print(response["messages"][-1].content)
+            print(result["messages"][-1].content)
         else:
-            # Regular message to agent
+            # обычный запрос к LLM
-            response = await agent.ainvoke(
+            result = await agent.ainvoke(
                 {"messages": [HumanMessage(content=user_input)]},
-                {"configurable": {"thread_id": "interactive"}},
+                {"configurable": {"thread_id": thread_id}},
             )
-            print(response["messages"][-1].content)
+            print(result["messages"][-1].content)
 
 # ---------------------------------------------------------------------------
-# Main entry point
+# Основной запуск
 # ---------------------------------------------------------------------------
 async def main():
-    # Optional: load initial documents
+    # Загрузим документы из каталога docs при старте
-    # await load_documents_from_dir("./initial_docs")
+    await load_documents_from_dir(DOCS_DIR)
     await interactive_loop()
 
 if __name__ == "__main__":