Агент с RAG‑памятью: rag_agent.py

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

@@ -1,4 +1,4 @@
-import os
+<|channel|>final code<|message|>import os
 from pathlib import Path
 from typing import List
 
@@ -6,169 +6,135 @@ 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, Tool
+from langchain.agents import create_agent, AgentExecutor
 from langchain.schema import Document
 
 # --------------------------------------------------------------------------- #
-# Конфигурация подключения к Qdrant и Ollama
+# 1. Векторное хранилище на базе Qdrant
 # --------------------------------------------------------------------------- #
 
-QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
+class KnowledgeBase:
-QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333))
+    """
-COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection")
+    Обёртка над QdrantVectorStore с автоматическим эмбеддингом через Ollama.
+    """
 
-EMBEDDING_MODEL = "nomic-embed-text"
+    def __init__(self, collection_name: str = "rag_collection", host: str = "localhost", port: int = 6333):
-LLM_MODEL = "llama3"
+        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
-
-# --------------------------------------------------------------------------- #
-# Векторное хранилище
-# --------------------------------------------------------------------------- #
-
-class RagVectorStore:
-    """Обёртка над QdrantVectorStore с автоматическим эмбеддингом Ollama."""
-
-    def __init__(self, collection_name: str = COLLECTION_NAME):
-        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
         self.store = QdrantVectorStore(
-            url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
-            collection_name=collection_name,
             embeddings=self.embeddings,
+            url=f"http://{host}:{port}",
+            collection_name=collection_name,
         )
-        # Создаём коллекцию, если её ещё нет
-        if not self.store.client.get_collection(collection_name).exists:
-            self.store.create_collection()
 
     def add_documents(self, docs: List[Document]) -> None:
-        """Добавляет список документов в хранилище."""
+        """Добавляет список документов в коллекцию."""
         self.store.add_documents(docs)
 
-    def search(
+    def search(self, query: str, limit: int = 5) -> List[Document]:
-        self,
+        """Поиск по запросу с ограничением на количество результатов."""
-        query: str,
+        return self.store.similarity_search(query=query, k=limit)
-        k: int = 5,
+
-    ) -> List[tuple[str, float]]:
-        """Семантический поиск по запросу. Возвращает кортежи (текст, релевантность)."""
-        results = self.store.similarity_search_with_score(query, k=k)
-        return [(doc.page_content, score) for doc, score in results]
 
 # --------------------------------------------------------------------------- #
-# Чанкинг документов
+# 2. Чанкинг документов
 # --------------------------------------------------------------------------- #
 
-def chunk_document(text: str, title: str) -> List[Document]:
+def chunk_text(text: str, title: str = "") -> List[Document]:
-    """Разбивает текст на чанки и возвращает список Document."""
+    """
+    Делит текст на чанки и возвращает список Document с метаданными.
+    """
     splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
+        chunk_size=500,
-        chunk_overlap=200,
+        chunk_overlap=50,
+        length_function=len
     )
     chunks = splitter.split_text(text)
-    return [
+    docs = [
         Document(page_content=chunk, metadata={"title": title})
         for chunk in chunks
     ]
+    return docs
+
 
 # --------------------------------------------------------------------------- #
-# Инструменты для агента
+# 3. Инструменты для агента
 # --------------------------------------------------------------------------- #
 
+kb = KnowledgeBase()
+
 @tool("search_knowledge_base")
-def search_knowledge_base(
+def search_knowledge_base(query: str, max_results: int = 5) -> List[str]:
-    query: str,
-    max_results: int = 5,
-) -> List[str]:
     """
     Семантический поиск в базе знаний.
-
+    Возвращает список строк с содержимым найденных чанков.
-    Parameters:
-        query (str): поисковый запрос.
-        max_results (int): максимальное количество результатов.
-
-    Returns:
-        list of strings: найденные тексты с оценкой релевантности.
     """
-    results = rag_store.search(query, k=max_results)
+    results = kb.search(query=query, limit=max_results)
-    return [f"{text} (score={score:.4f})" for text, score in results]
+    return [doc.page_content for doc in results]
+
 
 @tool("add_to_knowledge_base")
-def add_to_knowledge_base(
+def add_to_knowledge_base(content: str, title: str) -> str:
-    content: str,
-    title: str,
-) -> str:
     """
-    Добавляет новый документ в базу знаний.
+    Добавляет новый документ в базу знаний после чанкинга.
-
+    Возвращает подтверждение количества добавленных чанков.
-    Parameters:
-        content (str): полный текст документа.
-        title (str): заголовок/название документа.
-
-    Returns:
-        str: сообщение об успехе.
     """
-    docs = chunk_document(content, title)
+    docs = chunk_text(content, title=title)
-    rag_store.add_documents(docs)
+    kb.add_documents(docs)
-    return f"Документ '{title}' успешно добавлен в базу знаний."
+    return f"Добавлено {len(docs)} чанков из документа '{title}'."
 
 # --------------------------------------------------------------------------- #
-# Инициализация хранилища
+# 4. Создание агента
 # --------------------------------------------------------------------------- #
 
-rag_store = RagVectorStore()
+llm = OllamaLLM(model="llama3")
 
-# --------------------------------------------------------------------------- #
+agent_executor: AgentExecutor
-# Создание агента с RAG‑интеграцией
-# --------------------------------------------------------------------------- #
 
-llm = OllamaLLM(model=LLM_MODEL)
+def init_agent() -> None:
-
+    """
-tools: List[Tool] = [
+    Инициализирует агент с системным промптом, указывающим использовать RAG.
-    Tool.from_function(
+    """
-        func=search_knowledge_base,
+    global agent_executor
-        name="search_knowledge_base",
+    system_prompt = (
-        description="Используйте для поиска информации в базе знаний.",
+        "Ты AI-ассистент. При необходимости ищи информацию в базе знаний "
-    ),
+        "с помощью инструмента search_knowledge_base и добавляй новые данные "
-    Tool.from_function(
+        "через add_to_knowledge_base."
-        func=add_to_knowledge_base,
+    )
-        name="add_to_knowledge_base",
+    tools = [search_knowledge_base, add_to_knowledge_base]
-        description="Добавьте новый документ в базу знаний.",
+    agent_executor = create_agent(
-    ),
-]
-
-agent = create_agent(
         llm=llm,
         tools=tools,
-    system_message=(
+        system_prompt=system_prompt,
-        "Вы — интеллектуальный агент, способный искать и добавлять информацию "
+        verbose=True
-        "в локальную базу знаний. При необходимости используйте инструменты "
-        "'search_knowledge_base' и 'add_to_knowledge_base'."
-    ),
     )
 
-agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
-
 # --------------------------------------------------------------------------- #
-# Клиент для загрузки документов из директории
+# 5. Инициализация из директории (для загрузки файлов)
 # --------------------------------------------------------------------------- #
 
-def load_documents_from_directory(directory: str) -> None:
+def load_documents_from_dir(directory: str) -> None:
-    """Загружает все текстовые файлы из указанной папки в базу знаний."""
+    """
-    path = Path(directory)
+    Загружает все .txt файлы из указанной папки в базу знаний.
-    for file_path in path.rglob("*"):
+    """
-        if file_path.suffix.lower() in {".txt", ".md"}:
+    for file_path in Path(directory).glob("*.txt"):
-            text = file_path.read_text(encoding="utf-8")
+        with open(file_path, "r", encoding="utf-8") as f:
+            content = f.read()
         title = file_path.stem
-            add_to_knowledge_base(content=text, title=title)
+        add_to_knowledge_base(content=content, title=title)
 
 # --------------------------------------------------------------------------- #
-# Интерактивный клиент
+# 6. Интерактивный клиент
 # --------------------------------------------------------------------------- #
 
-def interactive_cli():
+def interactive_cli() -> None:
-    """Простой CLI для взаимодействия с агентом."""
+    """
-    print("=== RAG Агент ===")
+    Простая CLI с командами /add, /search и /quit.
-    print("Команды: /add <path>, /search <query>, /quit")
+    """
+    init_agent()
+    print("=== RAG-агент готов к работе ===")
     while True:
         try:
-            user_input = input("> ").strip()
+            user_input = input("\n> ").strip()
         except (EOFError, KeyboardInterrupt):
             break
 
@@ -176,31 +142,23 @@ def interactive_cli():
             continue
 
         if user_input.lower() == "/quit":
-            print("Выход.")
+            print("Завершение работы.")
             break
         elif user_input.startswith("/add"):
-            _, *path_parts = user_input.split(maxsplit=1)
+            # /add путь_к_файлу
-            if path_parts:
+            _, path = user_input.split(maxsplit=1)
-                file_path = Path(path_parts[0])
+            with open(path, "r", encoding="utf-8") as f:
-                if file_path.is_file():
+                content = f.read()
-                    content = file_path.read_text(encoding="utf-8")
+            title = Path(path).stem
-                    title = file_path.stem
             print(add_to_knowledge_base(content, title))
-                else:
-                    print(f"Файл не найден: {file_path}")
-            else:
-                print("Укажите путь к файлу.")
         elif user_input.startswith("/search"):
-            _, *query_parts = user_input.split(maxsplit=1)
+            # /search запрос
-            if query_parts:
+            _, query = user_input.split(maxsplit=1)
-                query = query_parts[0]
             results = search_knowledge_base(query, max_results=3)
-                for r in results:
+            for i, res in enumerate(results, 1):
-                    print(r)
+                print(f"\n--- Результат {i} ---\n{res}")
         else:
-                print("Укажите поисковый запрос.")
+            # Любой другой ввод считается запросом к агенту
-        else:
-            # Любой другой ввод передаём агенту
             response = agent_executor.invoke({"input": user_input})
             print(response["output"])