Обновить solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py

solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py

@@ -1,3 +1,4 @@
+# rag_agent.py
 import os
 from pathlib import Path
 from typing import List
@@ -6,161 +7,197 @@ from langchain_ollama import OllamaEmbeddings, OllamaLLM
 from langchain_qdrant import QdrantVectorStore
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain.tools import tool
-from langchain.agents import create_agent, AgentExecutor
+from langchain.agents import create_react_agent, AgentExecutor
+from langchain.prompts import PromptTemplate
 from langchain.schema import Document
-from rag_tools import search_knowledge_base, add_to_knowledge_base
+from qdrant_client import QdrantClient as _QdrantClient
+from qdrant_client.http.models import Distance, VectorParams
 
-# --------------------------------------------------------------------------- #
-# 1. Векторное хранилище на базе Qdrant
-# --------------------------------------------------------------------------- #
+# ---------------------------------------------------------------------------
+# Конфигурация
+# ---------------------------------------------------------------------------
 
-class KnowledgeBase:
-    """
-    Обёртка над QdrantVectorStore с автоматическим эмбеддингом через Ollama.
-    """
+QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
+QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
+COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection")
+EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "nomic-embed-text")
+LLM_MODEL = os.getenv("LLM_MODEL", "llama3")
+
+# ---------------------------------------------------------------------------
+# Векторное хранилище (Qdrant + Ollama-эмбеддинги)
+# ---------------------------------------------------------------------------
+
+class RagVectorStore:
+    """Обёртка над QdrantVectorStore с эмбеддингами через Ollama."""
+
+    def __init__(self, collection_name: str = COLLECTION_NAME):
+        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
+
+        # Низкоуровневый клиент для управления коллекцией
+        self._client = _QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
+
+        # Создаём коллекцию, если её ещё нет
+        if not self._collection_exists(collection_name):
+            self._client.create_collection(
+                collection_name=collection_name,
+                vectors_config=VectorParams(size=768, distance=Distance.COSINE),
+            )
 
-    def __init__(self, collection_name: str = "rag_collection", host: str = "localhost", port: int = 6333):
-        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
         self.store = QdrantVectorStore(
-            embeddings=self.embeddings,
-            url=f"http://{host}:{port}",
+            client=self._client,
             collection_name=collection_name,
+            embedding=self.embeddings,
         )
 
+    def _collection_exists(self, name: str) -> bool:
+        try:
+            self._client.get_collection(name)
+            return True
+        except Exception:
+            return False
+
     def add_documents(self, docs: List[Document]) -> None:
-        """Добавляет список документов в коллекцию."""
         self.store.add_documents(docs)
 
-    def search(self, query: str, limit: int = 5) -> List[Document]:
-        """Поиск по запросу с ограничением на количество результатов."""
-        return self.store.similarity_search(query=query, k=limit)
+    def search(self, query: str, k: int = 5) -> List[tuple]:
+        results = self.store.similarity_search_with_score(query, k=k)
+        return [(doc.page_content, score) for doc, score in results]
 
 
-# --------------------------------------------------------------------------- #
-# 2. Чанкинг документов
-# --------------------------------------------------------------------------- #
+# ---------------------------------------------------------------------------
+# Чанкинг
+# ---------------------------------------------------------------------------
 
-def chunk_text(text: str, title: str = "") -> List[Document]:
-    """
-    Делит текст на чанки и возвращает список Document с метаданными.
-    """
+def chunk_document(text: str, title: str) -> List[Document]:
+    """Разбивает текст на чанки с метаданными."""
     splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500,
-        chunk_overlap=50,
-        length_function=len
+        chunk_size=1000,
+        chunk_overlap=200,
+        separators=["\n\n", "\n", " ", ""],
     )
-    chunks = splitter.split_text(text)
-    docs = [
+    return [
         Document(page_content=chunk, metadata={"title": title})
-        for chunk in chunks
+        for chunk in splitter.split_text(text)
     ]
-    return docs
 
 
-# --------------------------------------------------------------------------- #
-# 3. Инструменты для агента
-# --------------------------------------------------------------------------- #
+# ---------------------------------------------------------------------------
+# Инициализация хранилища (глобальный синглтон для инструментов)
+# ---------------------------------------------------------------------------
 
-kb = KnowledgeBase()
+rag_store = RagVectorStore()
 
-@tool("search_knowledge_base")
+
+# ---------------------------------------------------------------------------
+# LangChain-инструменты агента (@tool)
+# ---------------------------------------------------------------------------
+
+@tool
 def search_knowledge_base(query: str, max_results: int = 5) -> List[str]:
     """
     Семантический поиск в базе знаний.
-    Возвращает список строк с содержимым найденных чанков.
+
+    Args:
+        query: поисковый запрос.
+        max_results: максимальное количество результатов.
+
+    Returns:
+        Список найденных фрагментов с оценкой релевантности.
     """
-    results = kb.search(query=query, limit=max_results)
-    return [doc.page_content for doc in results]
+    results = rag_store.search(query, k=max_results)
+    return [f"[score={score:.4f}] {text}" for text, score in results]
 
 
-@tool("add_to_knowledge_base")
+@tool
 def add_to_knowledge_base(content: str, title: str) -> str:
     """
-    Добавляет новый документ в базу знаний после чанкинга.
-    Возвращает подтверждение количества добавленных чанков.
+    Добавляет новый документ в базу знаний.
+
+    Args:
+        content: полный текст документа.
+        title: заголовок документа.
+
+    Returns:
+        Сообщение об успехе.
     """
-    docs = chunk_text(content, title=title)
-    kb.add_documents(docs)
-    return f"Добавлено {len(docs)} чанков из документа '{title}'."
+    docs = chunk_document(content, title)
+    rag_store.add_documents(docs)
+    return f"Документ '{title}' успешно добавлен в базу знаний."
 
-# --------------------------------------------------------------------------- #
-# 4. Создание агента
-# --------------------------------------------------------------------------- #
 
-LLM_MODEL = "llama3"
+# ---------------------------------------------------------------------------
+# Агент на базе LangChain (create_react_agent)
+# ---------------------------------------------------------------------------
+
 llm = OllamaLLM(model=LLM_MODEL)
 
 tools = [search_knowledge_base, add_to_knowledge_base]
 
 prompt = PromptTemplate.from_template(
-    "Вы — интеллектуальный агент с доступом к базе знаний.\n"
-    "Используйте инструменты 'search_knowledge_base' и 'add_to_knowledge_base' при необходимости.\n\n"
+    "Вы — интеллектуальный агент с доступом к локальной базе знаний.\n"
+    "При необходимости используйте инструменты search_knowledge_base и add_to_knowledge_base.\n\n"
     "Доступные инструменты:\n{tools}\n\n"
-    "Формат ответа:\n"
+    "Формат ответа (строго соблюдайте):\n"
+    "Question: вопрос пользователя\n"
+    "Thought: что нужно сделать\n"
+    "Action: инструмент из [{tool_names}]\n"
+    "Action Input: входные данные для инструмента\n"
+    "Observation: результат инструмента\n"
+    "... (можно повторять Thought/Action/Observation)\n"
+    "Thought: теперь я знаю ответ\n"
+    "Final Answer: итоговый ответ пользователю\n\n"
+    "Begin!\n\n"
     "Question: {input}\n"
-    "Thought: {agent_scratchpad}\n"
-    "Action: действие из [{tool_names}]\n"
-    "Action Input: входные данные\n"
-    "Observation: результат\n"
-    "Final Answer: итоговый ответ"
+    "Thought: {agent_scratchpad}"
 )
 
 agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
 agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
-# --------------------------------------------------------------------------- #
-# 5. Инициализация из директории (для загрузки файлов)
-# --------------------------------------------------------------------------- #
 
-def load_documents_from_dir(directory: str) -> None:
-    """
-    Загружает все .txt файлы из указанной папки в базу знаний.
-    """
-    for file_path in Path(directory).glob("*.txt"):
-        with open(file_path, "r", encoding="utf-8") as f:
-            content = f.read()
-        title = file_path.stem
-        add_to_knowledge_base(content=content, title=title)
+# ---------------------------------------------------------------------------
+# Загрузка документов из директории (для init_loader.py)
+# ---------------------------------------------------------------------------
 
-# --------------------------------------------------------------------------- #
-# 6. Интерактивный клиент
-# --------------------------------------------------------------------------- #
+def load_documents_from_directory(directory: str) -> None:
+    """Загружает все .txt и .md файлы из папки в базу знаний."""
+    path = Path(directory)
+    for file_path in path.rglob("*"):
+        if file_path.suffix.lower() in {".txt", ".md"}:
+            text = file_path.read_text(encoding="utf-8")
+            add_to_knowledge_base.invoke({"content": text, "title": file_path.stem})
+            print(f"Загружен: {file_path.name}")
 
-def interactive_cli() -> None:
-    """
-    Простая CLI с командами /add, /search и /quit.
-    """
-    init_agent()
-    print("=== RAG-агент готов к работе ===")
+
+# ---------------------------------------------------------------------------
+# Запуск напрямую — интерактивный CLI
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    print("=== RAG Агент ===")
+    print("Команды: /add <path>, /search <query>, /quit")
     while True:
         try:
-            user_input = input("\n> ").strip()
+            user_input = input("> ").strip()
         except (EOFError, KeyboardInterrupt):
             break
 
         if not user_input:
             continue
-
         if user_input.lower() == "/quit":
-            print("Завершение работы.")
+            print("Выход.")
             break
-        elif user_input.startswith("/add"):
-            # /add путь_к_файлу
-            _, path = user_input.split(maxsplit=1)
-            with open(path, "r", encoding="utf-8") as f:
-                content = f.read()
-            title = Path(path).stem
-            print(add_to_knowledge_base(content, title))
-        elif user_input.startswith("/search"):
-            # /search запрос
-            _, query = user_input.split(maxsplit=1)
-            results = search_knowledge_base(query, max_results=3)
-            for i, res in enumerate(results, 1):
-                print(f"\n--- Результат {i} ---\n{res}")
+        elif user_input.startswith("/add "):
+            file_path = Path(user_input[5:].strip())
+            if file_path.is_file():
+                text = file_path.read_text(encoding="utf-8")
+                print(add_to_knowledge_base.invoke({"content": text, "title": file_path.stem}))
+            else:
+                print(f"Файл не найден: {file_path}")
+        elif user_input.startswith("/search "):
+            query = user_input[8:].strip()
+            for r in search_knowledge_base.invoke({"query": query, "max_results": 3}):
+                print(r)
         else:
-            # Любой другой ввод считается запросом к агенту
             response = agent_executor.invoke({"input": user_input})
-            print(response["output"])
-
-if __name__ == "__main__":
-    interactive_cli()
+            print(response["output"])