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

main.py

@@ -1,58 +1,152 @@
 import os
+import sys
 import asyncio
-from dotenv import load_dotenv
-from langchain_openai import ChatOpenAI
+from pathlib import Path
+from typing import List
+
+from langchain_ollama import Ollama, OllamaEmbeddings
+from langchain_core.documents import Document
+from langchain_qdrant import QdrantVectorStore
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.tools import tool
 from langchain_core.messages import HumanMessage
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-from tools import search_knowledge_base, add_to_knowledge_base
+from qdrant_client import QdrantClient
 
-load_dotenv()
+# Инициализация эмбеддингов и LLM через Ollama
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+llm = Ollama(model="llama3")
 
-llm = ChatOpenAI(
-    model="openai/gpt-oss-20b:free",
-    base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
-    temperature=0.0,
+# Инициализация Qdrant
+client = QdrantClient(host="localhost", port=6333)
+collection_name = "knowledge"
+vector_store = QdrantVectorStore(
+    client=client,
+    collection_name=collection_name,
+    embedding=embeddings,
 )
 
-backend = CompositeBackend([
+# Чанкинг
+splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+
+# Инструмент: поиск в базе знаний
+@tool
+def search_knowledge_base(query: str, max_results: int = 3) -> str:
+    """Search the knowledge base for relevant information."""
+    docs: List[Document] = vector_store.similarity_search(query, k=max_results)
+    if not docs:
+        return "No results."
+    return "\n".join(doc.page_content for doc in docs)
+
+# Инструмент: добавление документа в базу знаний
+@tool
+def add_to_knowledge_base(content: str, title: str = "doc") -> str:
+    """Add content to the knowledge base."""
+    chunks = splitter.split_text(content)
+    docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks]
+    vector_store.add_documents(docs)
+    return f"Added: {title} ({len(chunks)} chunks)."
+
+# Backend для deepagents
+backend = CompositeBackend(
+    [
         LocalShellBackend(workspace_dir="./workspace"),
         FilesystemBackend(),
-])
+    ]
+)
+
+system_prompt = (
+    "You are a helpful agent with access to a knowledge base. "
+    "Use the tools to search and add information. "
+    "When you need to retrieve information, call search_knowledge_base. "
+    "When you need to store new information, call add_to_knowledge_base."
+)
 
 agent = create_deep_agent(
     model=llm,
     tools=[search_knowledge_base, add_to_knowledge_base],
     backend=backend,
-    system_prompt="You are a helpful knowledge assistant. Use the tools to search and add documents.",
+    system_prompt=system_prompt,
 )
 
-async def interactive_loop():
-    thread_id = "interactive-session"
-    print("Welcome to RAG Agent. Commands: /add <title> <content>, /search <query>, /quit")
+# Загрузка документов из директории
+def load_documents(dir_path: str) -> None:
+    """Load all .txt files from dir_path into the knowledge base."""
+    path = Path(dir_path)
+    if not path.is_dir():
+        print(f"Directory not found: {dir_path}")
+        return
+    for file_path in path.glob("*.txt"):
+        try:
+            content = file_path.read_text(encoding="utf-8")
+            title = file_path.stem
+            result = add_to_knowledge_base(content, title)
+            print(f"Loaded {file_path.name}: {result}")
+        except Exception as e:
+            print(f"Error loading {file_path.name}: {e}")
+
+# Интерактивный клиент
+async def run_cli() -> None:
+    thread_id = "session-1"
+    print("Welcome to the RAG agent CLI.")
+    print("Commands:")
+    print("  /add    - add a new document")
+    print("  /search - search the knowledge base")
+    print("  /quit   - exit")
     while True:
-        user_input = input(">> ")
-        if user_input.strip() == "/quit":
+        try:
+            user_input = input("\n> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nExiting.")
+            break
+
+        if not user_input:
+            continue
+
+        if user_input.lower() == "/quit":
             print("Goodbye.")
             break
-        if user_input.startswith("/add"):
-            try:
-                _, title, content = user_input.split(" ", 2)
-            except ValueError:
-                print("Usage: /add <title> <content>")
+
+        if user_input.lower() == "/add":
+            title = input("Title: ").strip()
+            print("Enter content (end with a single line containing only 'END'):")
+            lines: List[str] = []
+            while True:
+                line = input()
+                if line.strip() == "END":
+                    break
+                lines.append(line)
+            content = "\n".join(lines)
+            result = add_to_knowledge_base(content, title)
+            print(result)
             continue
-            message = HumanMessage(content=f"Add document titled '{title}' with content: {content}")
-        elif user_input.startswith("/search"):
-            query = user_input[len("/search"):].strip()
-            message = HumanMessage(content=f"Search knowledge base for: {query}")
-        else:
-            message = HumanMessage(content=user_input)
-        result = await agent.ainvoke(
-            {"messages": [message]},
+
+        if user_input.lower() == "/search":
+            query = input("Enter search query: ").strip()
+            if not query:
+                print("Empty query.")
+                continue
+            result = search_knowledge_base(query)
+            print("\nSearch results:")
+            print(result)
+            continue
+
+        # Any other message is sent to the agent
+        response = await agent.ainvoke(
+            {"messages": [HumanMessage(content=user_input)]},
             {"configurable": {"thread_id": thread_id}},
         )
-        print(result["messages"][-1].content)
+        agent_reply = response["messages"][-1].content
+        print(f"\nAgent: {agent_reply}")
+
+def main() -> None:
+    # Optional: load documents from a directory passed as first argument
+    if len(sys.argv) > 1:
+        dir_path = sys.argv[1]
+        print(f"Loading documents from {dir_path}...")
+        load_documents(dir_path)
+    asyncio.run(run_cli())
 
 if __name__ == "__main__":
-    asyncio.run(interactive_loop())
+    main()