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

main.py

@@ -1,119 +1,77 @@
-import os
 import asyncio
-from langchain_ollama import Ollama, OllamaEmbeddings
+import os
+from pathlib import Path
+
+from langchain_core.messages import HumanMessage
 from langchain.tools import tool
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-from langchain_qdrant import QdrantVectorStore
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+from utils import llm, vector_store, splitter
 from langchain_core.documents import Document
-from langchain_core.messages import HumanMessage
 
-# Инициализация эмбеддинговой модели Ollama
-embeddings = OllamaEmbeddings(model="nomic-embed-text")
-
-# Инициализация векторного хранилища Qdrant
-vector_store = QdrantVectorStore(
-    url="http://localhost:6333",
-    collection_name="knowledge",
-    embedding_function=embeddings
-)
-
-# Инструмент: поиск в базе знаний
 @tool
 def search_knowledge_base(query: str, max_results: int = 3) -> str:
     """Search the knowledge base for relevant information."""
     docs = vector_store.similarity_search(query, k=max_results)
     return "\n".join(d.page_content for d in docs) if docs else "No results."
 
-# Инструмент: добавление документа в базу знаний
 @tool
 def add_to_knowledge_base(content: str, title: str = "doc") -> str:
     """Add content to the knowledge base."""
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     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)"
+    return f"Added {len(docs)} chunks for {title}."
 
-# Инициализация LLM Ollama
-llm = Ollama(model="llama3")
-
-# Backend для deepagents
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
-# Системный промпт агента
-system_prompt = (
-    "You are a helpful agent with access to a knowledge base. "
-    "Use the provided tools to search and add information. "
-    "When answering user queries, first search the knowledge base and then provide a concise response."
-)
-
-# Создание агента
 agent = create_deep_agent(
     model=llm,
     tools=[search_knowledge_base, add_to_knowledge_base],
     backend=backend,
-    system_prompt=system_prompt,
+    system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.",
 )
 
-# Загрузка документов из директории в векторное хранилище
-def load_documents_from_dir(dir_path: str):
-    for root, dirs, files in os.walk(dir_path):
-        for file in files:
-            if file.lower().endswith(".txt"):
-                path = os.path.join(root, file)
-                with open(path, "r", encoding="utf-8") as f:
-                    content = f.read()
-                title = os.path.splitext(file)[0]
-                result = add_to_knowledge_base(content, title)
-                print(result)
-
-# Интерактивный клиент
-async def interactive_cli():
-    print("Welcome to the RAG agent. Commands: /add <file_path>, /search <query>, /quit")
+async def main():
+    print("RAG Agent CLI. Commands: /add, /search, /quit")
     while True:
-        user_input = input(">> ")
-        if not user_input.strip():
+        user_input = input(">> ").strip()
+        if not user_input:
             continue
-        if user_input.startswith("/add"):
-            parts = user_input.split(maxsplit=1)
-            if len(parts) < 2:
-                print("Usage: /add <file_path>")
-                continue
-            file_path = parts[1]
-            if not os.path.isfile(file_path):
-                print(f"File not found: {file_path}")
-                continue
-            with open(file_path, "r", encoding="utf-8") as f:
-                content = f.read()
-            title = os.path.splitext(os.path.basename(file_path))[0]
+        if user_input.lower() == "/quit":
+            print("Goodbye.")
+            break
+        if user_input.lower().startswith("/add"):
+            title = input("Title: ").strip()
+            print("Enter content (end with a single line containing only END):")
+            lines = []
+            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)
-        elif user_input.startswith("/search"):
-            parts = user_input.split(maxsplit=1)
-            if len(parts) < 2:
-                print("Usage: /search <query>")
-                continue
-            query = parts[1]
-            response = await agent.ainvoke(
-                {"messages": [HumanMessage(content=query)]},
-                {"configurable": {"thread_id": "session-1"}},
-            )
-            print(response["messages"][-1].content)
-        elif user_input.startswith("/quit"):
-            print("Goodbye!")
-            break
-        else:
-            print("Unknown command. Use /add, /search, /quit")
-
-async def main():
-    # При необходимости загрузить начальные документы
-    # load_documents_from_dir("./docs")
-    await interactive_cli()
+            continue
+        if user_input.lower().startswith("/search"):
+            query = input("Query: ").strip()
+            max_results_str = input("Max results (default 3): ").strip()
+            max_results = int(max_results_str) if max_results_str.isdigit() else 3
+            result = search_knowledge_base(query, max_results)
+            print("Search results:")
+            print(result)
+            continue
+        # Regular message to agent
+        response = await agent.ainvoke(
+            {"messages": [HumanMessage(content=user_input)]},
+            {"configurable": {"thread_id": "session-1"}},
+        )
+        print(response["messages"][-1].content)
 
 if __name__ == "__main__":
     asyncio.run(main())