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

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

@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+solution.py – Агент с RAG‑памятью на базе Qdrant и Ollama.
+"""
+
+import os
+import sys
+import argparse
+from pathlib import Path
+from typing import List, Dict
+
+from langchain_ollama import ChatOllama, OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.tools import tool
+from langchain.agents import create_agent, AgentExecutor, Tool
+from langchain.schema import Document
+
+# --------------------------------------------------------------------------- #
+# Конфигурация
+# --------------------------------------------------------------------------- #
+
+QDRANT_HOST = "localhost"
+QDRANT_PORT = 6333
+COLLECTION_NAME = "knowledge_base"
+
+EMBEDDING_MODEL = "nomic-embed-text"   # Ollama embedding model
+LLM_MODEL = "llama3"                   # Ollama LLM
+
+# --------------------------------------------------------------------------- #
+# Векторное хранилище и вспомогательные функции
+# --------------------------------------------------------------------------- #
+
+class KnowledgeBase:
+    """
+    Обёртка над QdrantVectorStore с поддержкой чанкинга и инструментов.
+    """
+
+    def __init__(self,
+                 host: str = QDRANT_HOST,
+                 port: int = QDRANT_PORT,
+                 collection_name: str = COLLECTION_NAME):
+        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
+        # Создаём/подключаем коллекцию
+        self.store = QdrantVectorStore(
+            embedding=self.embeddings,
+            url=f"http://{host}:{port}",
+            collection_name=collection_name,
+        )
+        # Чанкер для разбивки документов
+        self.splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+
+    def add_document(self, content: str, title: str) -> None:
+        """
+        Добавляет документ в базу после разбиения на чанки.
+        """
+        # Разбиваем текст
+        chunks = self.splitter.split_text(content)
+        docs = [
+            Document(page_content=chunk,
+                     metadata={"title": title, "source": f"{title} (chunk {i})"})
+            for i, chunk in enumerate(chunks)
+        ]
+        # Добавляем в Qdrant
+        self.store.add_documents(docs)
+
+    def search(self, query: str, max_results: int = 5) -> List[Dict]:
+        """
+        Семантический поиск по базе знаний.
+        Возвращает список словарей с полями 'title', 'source' и 'content'.
+        """
+        results = self.store.similarity_search_with_score(query, k=max_results)
+        output = []
+        for doc, score in results:
+            out = {
+                "score": score,
+                "title": doc.metadata.get("title", ""),
+                "source": doc.metadata.get("source", ""),
+                "content": doc.page_content
+            }
+            output.append(out)
+        return output
+
+
+# --------------------------------------------------------------------------- #
+# Инструменты для агента
+# --------------------------------------------------------------------------- #
+
+kb = KnowledgeBase()
+
+@tool
+def search_knowledge_base(query: str, max_results: int) -> str:
+    """
+    Поиск в базе знаний.
+    Возвращает строку с найденными результатами в формате Markdown.
+    """
+    results = kb.search(query, max_results)
+    if not results:
+        return "Ничего не найдено."
+    lines = []
+    for r in results:
+        lines.append(f"**{r['title']}** ({r['source']}) – {r['score']:.2f}\n> {r['content'][:200]}...")
+    return "\n\n".join(lines)
+
+
+@tool
+def add_to_knowledge_base(content: str, title: str) -> str:
+    """
+    Добавление нового документа в базу знаний.
+    """
+    kb.add_document(content, title)
+    return f"Документ «{title}» успешно добавлен."
+
+
+# --------------------------------------------------------------------------- #
+# Создание агента
+# --------------------------------------------------------------------------- #
+
+def build_agent() -> AgentExecutor:
+    """
+    Возвращает готовый агент с инструментами RAG‑памяти.
+    """
+    system_prompt = (
+        "Ты – интеллектуальный ассистент, который использует локальную базу знаний. "
+        "При ответе на запросы сначала ищи релевантную информацию в базе через tool "
+        "`search_knowledge_base`. Если нужна новая информация, добавляй её через "
+        "`add_to_knowledge_base`."
+    )
+    tools = [
+        Tool.from_function(
+            func=search_knowledge_base,
+            name="search_knowledge_base",
+            description="Семантический поиск в базе знаний. Вход: query, max_results. Выход: Markdown с результатами."
+        ),
+        Tool.from_function(
+            func=add_to_knowledge_base,
+            name="add_to_knowledge_base",
+            description="Добавить новый документ в базу знаний. Вход: content, title. Выход: подтверждение."
+        )
+    ]
+    llm = ChatOllama(model=LLM_MODEL)
+    agent = create_agent(
+        llm=llm,
+        tools=tools,
+        system_prompt=system_prompt,
+        verbose=True
+    )
+    return AgentExecutor(agent=agent, tools=tools, verbose=True)
+
+
+# --------------------------------------------------------------------------- #
+# Загрузка документов из директории
+# --------------------------------------------------------------------------- #
+
+def load_documents_from_dir(directory: Path) -> None:
+    """
+    Загружает все текстовые файлы из указанной папки в базу знаний.
+    """
+    for file_path in directory.rglob("*"):
+        if file_path.suffix.lower() not in {".txt", ".md"}:
+            continue
+        content = file_path.read_text(encoding="utf-8")
+        title = file_path.stem
+        kb.add_document(content, title)
+        print(f"Загружено: {file_path}")
+
+
+# --------------------------------------------------------------------------- #
+# CLI клиент
+# --------------------------------------------------------------------------- #
+
+def interactive_cli(agent_executor: AgentExecutor) -> None:
+    """
+    Простая REPL‑интерфейс с командами /add, /search и /quit.
+    Любой другой ввод считается запросом к агенту.
+    """
+    print("=== Агент RAG ===")
+    print("Команды:")
+    print("/add <title> – добавить новый документ (будет запрошен контент)")
+    print("/search <query> – поиск в базе знаний")
+    print("/quit – выйти")
+
+    while True:
+        try:
+            inp = input("\n> ").strip()
+        except EOFError:
+            break
+        if not inp:
+            continue
+
+        if inp.lower() == "/quit":
+            break
+
+        if inp.startswith("/add "):
+            title = inp[5:].strip()
+            print("Введите содержимое документа (конец строки с EOF Ctrl+D):")
+            try:
+                content_lines = []
+                while True:
+                    line = input()
+                    content_lines.append(line)
+            except EOFError:
+                pass
+            content = "\n".join(content_lines)
+            result = add_to_knowledge_base(content, title)
+            print(result)
+
+        elif inp.startswith("/search "):
+            query = inp[8:].strip()
+            results = kb.search(query, max_results=5)
+            if not results:
+                print("Ничего не найдено.")
+            else:
+                for r in results:
+                    print(f"\n**{r['title']}** ({r['source']}) – {r['score']:.2f}")
+                    print(r["content"][:400] + "...")
+        else:
+            # Любой другой ввод – запрос к агенту
+            response = agent_executor.invoke({"input": inp})
+            print(response.get("output", ""))
+
+
+# --------------------------------------------------------------------------- #
+# Точка входа
+# --------------------------------------------------------------------------- #
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Агент с RAG‑памятью")
+    parser.add_argument(
+        "--load-dir",
+        type=str,
+        help="Путь к директории с документами для предварительной загрузки"
+    )
+    args = parser.parse_args()
+
+    if args.load_dir:
+        load_documents_from_dir(Path(args.load_dir))
+
+    agent_executor = build_agent()
+    interactive_cli(agent_executor)
+
+
+if __name__ == "__main__":
+    main()