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

2026-05-26 09:21:07 +00:00
parent 055ce3d080
commit 64ae908897
1 changed files with 187 additions and 0 deletions
@@ -0,0 +1,187 @@
+<|channel|>final code<|message|>import os
+from pathlib import Path
+from typing import List
+
+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.schema import Document
+
+# --------------------------------------------------------------------------- #
+# Конфигурация и клиент Qdrant
+# --------------------------------------------------------------------------- #
+
+QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
+QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333))
+COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection")
+
+# Инициализируем клиент Qdrant
+qdrant_store = QdrantVectorStore(
+    url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
+    collection_name=COLLECTION_NAME,
+    embeddings=OllamaEmbeddings(model="nomic-embed-text"),
+)
+
+# --------------------------------------------------------------------------- #
+# Чанкинг текста
+# --------------------------------------------------------------------------- #
+
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=500,  # символов в чанке
+    chunk_overlap=50,
+)
+
+# --------------------------------------------------------------------------- #
+# Инструменты агента
+# --------------------------------------------------------------------------- #
+
+@tool("search_knowledge_base")
+def search_knowledge_base(query: str, max_results: int = 5) -> List[Document]:
+    """
+    Семантический поиск по базе знаний.
+
+    :param query: поисковый запрос
+    :param max_results: максимальное количество результатов
+    :return: список документов с метаданными
+    """
+    results = qdrant_store.similarity_search(query, k=max_results)
+    return results
+
+
+@tool("add_to_knowledge_base")
+def add_to_knowledge_base(content: str, title: str) -> str:
+    """
+    Добавление нового документа в базу знаний.
+
+    :param content: полный текст документа
+    :param title: заголовок/название документа
+    :return: подтверждение добавления
+    """
+    # Разбиваем на чанки и сохраняем с метаданными
+    chunks = text_splitter.split_text(content)
+    docs = [
+        Document(page_content=chunk, metadata={"title": title})
+        for chunk in chunks
+    ]
+    qdrant_store.add_documents(docs)
+    return f"Документ '{title}' успешно добавлен в базу знаний."
+
+
+# --------------------------------------------------------------------------- #
+# Создание агента
+# --------------------------------------------------------------------------- #
+
+llm = OllamaLLM(model="llama3")
+
+tools: List[Tool] = [
+    Tool.from_function(
+        func=search_knowledge_base,
+        name="search_knowledge_base",
+        description="Используйте для поиска информации в базе знаний. Входные параметры: query, max_results.",
+    ),
+    Tool.from_function(
+        func=add_to_knowledge_base,
+        name="add_to_knowledge_base",
+        description="Добавляет новый документ в базу знаний. Параметры: content, title.",
+    ),
+]
+
+agent = create_agent(
+    llm=llm,
+    tools=tools,
+    system_message=(
+        "Вы — интеллектуальный агент с доступом к базе знаний. "
+        "При необходимости используйте инструмент search_knowledge_base для поиска информации. "
+        "Если нужно добавить новый материал, применяйте add_to_knowledge_base."
+    ),
+)
+
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+
+
+# --------------------------------------------------------------------------- #
+# Загрузка документов из директории
+# --------------------------------------------------------------------------- #
+
+def load_documents_from_directory(directory: str) -> None:
+    """
+    Загружает все текстовые файлы из указанной папки в базу знаний.
+
+    :param directory: путь к каталогу с документами
+    """
+    path = Path(directory)
+    for file_path in path.rglob("*"):
+        if file_path.suffix.lower() not in {".txt", ".md"}:
+            continue
+        content = file_path.read_text(encoding="utf-8")
+        title = file_path.stem
+        add_to_knowledge_base(content, title)
+
+
+# --------------------------------------------------------------------------- #
+# Интерактивный клиент
+# --------------------------------------------------------------------------- #
+
+def interactive_cli():
+    """
+    Простая CLI для взаимодействия с агентом.
+    Команды:
+      /add <path>   - добавить документ из файла
+      /search <q>   - поиск по базе знаний
+      /quit         - выйти
+      любой другой ввод передаётся агенту как вопрос
+    """
+    print("=== RAG Агент ===")
+    while True:
+        try:
+            user_input = input("\n> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            break
+
+        if not user_input:
+            continue
+
+        if user_input.lower() == "/quit":
+            print("До свидания!")
+            break
+        elif user_input.startswith("/add"):
+            _, *parts = user_input.split(maxsplit=1)
+            if parts:
+                file_path = parts[0]
+                try:
+                    content = Path(file_path).read_text(encoding="utf-8")
+                    title = Path(file_path).stem
+                    print(add_to_knowledge_base(content, title))
+                except Exception as e:
+                    print(f"Ошибка при добавлении: {e}")
+            else:
+                print("Используйте /add <путь_к_файлу>")
+        elif user_input.startswith("/search"):
+            _, *parts = user_input.split(maxsplit=1)
+            if parts:
+                query = parts[0]
+                results = search_knowledge_base(query, max_results=3)
+                for i, doc in enumerate(results, 1):
+                    print(f"\nРезультат {i}:")
+                    print(f"Титул: {doc.metadata.get('title', 'Неизвестно')}")
+                    print(doc.page_content[:500], "...")
+            else:
+                print("Используйте /search <запрос>")
+        else:
+            # Любой другой ввод считается вопросом агента
+            response = agent_executor.run(user_input)
+            print(response)
+
+
+# --------------------------------------------------------------------------- #
+# Точка входа
+# --------------------------------------------------------------------------- #
+
+if __name__ == "__main__":
+    # При запуске можно загрузить документы из папки docs/
+    docs_dir = os.getenv("DOCS_DIR", "docs")
+    if Path(docs_dir).exists():
+        load_documents_from_directory(docs_dir)
+    interactive_cli()