89 lines
4.5 KiB
Python
89 lines
4.5 KiB
Python
# 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)}" |