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

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

@@ -0,0 +1,89 @@
+# tools.py
+"""
+Модуль с инструментами для агента.
+В файле определены два инструмента:
+    1. search_knowledge_base – семантический поиск в Qdrant.
+    2. add_to_knowledge_base   – добавление документа (с чанками) в базу.
+
+Инструменты реализованы через декоратор @tool из langchain.tools,
+что позволяет автоматически использовать их в агентах LangChain.
+"""
+
+from pathlib import Path
+from typing import List, Tuple
+
+# ────────────────────────  Зависимости  ────────────────────────
+from langchain.embeddings.ollama import OllamaEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores.qdrant import QdrantVectorStore
+from langchain.tools import tool
+
+# ────────────────────────  Конфигурация  ────────────────────────
+# Путь к локальной базе Qdrant (предполагается, что она уже запущена)
+QDRANT_URL = "http://localhost:6333"
+COLLECTION_NAME = "knowledge_base"
+
+# Создаём один экземпляр векторного хранилища и эмбеддеров,
+# чтобы не создавать их каждый раз при вызове инструментов.
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+vectorstore = QdrantVectorStore(
+    client_kwargs={"url": QDRANT_URL},
+    collection_name=COLLECTION_NAME,
+    embedding_function=embeddings.embed_query,  # функция для получения эмбеддингов
+)
+
+# ────────────────────────  Инструмент: поиск  ────────────────────────
+@tool("search_knowledge_base",
+      description="Семантический поиск в базе знаний. "
+                  "Возвращает список строк с найденными документами и их релевантностью.")
+def search_knowledge_base(query: str, max_results: int = 5) -> List[Tuple[str, float]]:
+    """
+    Выполняет семантический поиск по запросу.
+
+    Args:
+        query (str): поисковый запрос.
+        max_results (int): максимальное количество результатов.
+
+    Returns:
+        List[Tuple[str, float]]: список кортежей (текст документа, релевантность).
+    """
+    # similarity_search_with_score возвращает [(Document, score), ...]
+    results = vectorstore.similarity_search_with_score(query, k=max_results)
+    return [(doc.page_content, score) for doc, score in results]
+
+
+# ────────────────────────  Инструмент: добавление  ────────────────────────
+@tool("add_to_knowledge_base",
+      description="Добавляет новый документ в базу знаний. "
+                  "Документ разбивается на чанки и сохраняется с метаданными.")
+def add_to_knowledge_base(content: str, title: str) -> str:
+    """
+    Добавляет текстовый контент в Qdrant.
+
+    Args:
+        content (str): полный текст документа.
+        title  (str): заголовок/имя документа.
+
+    Returns:
+        str: сообщение об успешном добавлении и количестве чанков.
+    """
+    # Разбиваем документ на чанки
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+        separators=["\n\n", "\n", " ", ""],
+    )
+    chunks = splitter.split_text(content)
+
+    # Создаём объекты Document с метаданными
+    from langchain.schema import Document
+
+    documents: List[Document] = [
+        Document(page_content=chunk, metadata={"title": title})
+        for chunk in chunks
+    ]
+
+    # Добавляем векторные представления в Qdrant
+    vectorstore.add_documents(documents)
+
+    return f"Документ '{title}' добавлен. Количество чанков: {len(chunks)}"