Агент с RAG‑памятью: qdrant_client.py
This commit is contained in:
@@ -1,154 +1,175 @@
|
||||
# qdrant_client.py
|
||||
"""
|
||||
Модуль инициализации клиента Qdrant, добавления документов и поиска по эмбеддингам Ollama.
|
||||
Использует:
|
||||
- langchain_qdrant.QdrantVectorStore для работы с коллекцией Qdrant
|
||||
- langchain_ollama.OllamaEmbeddings для генерации векторных представлений
|
||||
- RecursiveCharacterTextSplitter для разбиения больших текстов на чанки
|
||||
Модуль инициализации клиента Qdrant и простых операций с ним.
|
||||
Используется в агенте для RAG‑памяти.
|
||||
|
||||
Функциональность:
|
||||
1. Инициализация клиента и создание/получение коллекции.
|
||||
2. Добавление документа (с заголовком) в базу знаний с автоматическим чанкингом.
|
||||
3. Поиск по семантическому запросу с ограничением количества результатов.
|
||||
Поддержка:
|
||||
- Создание/получение коллекции
|
||||
- Добавление документов (с чанками) с эмбеддингами Ollama
|
||||
- Семантический поиск по запросу
|
||||
|
||||
Author: ChatGPT
|
||||
Требуемые зависимости:
|
||||
pip install langchain langchain-qdrant langchain-ollama qdrant-client
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Any
|
||||
from typing import Iterable, List, Tuple
|
||||
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
from langchain_qdrant import QdrantVectorStore
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from qdrant_client import QdrantClient as _QdrantClient
|
||||
from qdrant_client.http.models import Distance, VectorParams
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Конфигурация (можно переопределить через переменные окружения)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
QDRANT_HOST: str = os.getenv("QDRANT_HOST", "localhost")
|
||||
QDRANT_PORT: int = int(os.getenv("QDRANT_PORT", "6333"))
|
||||
COLLECTION_NAME: str = os.getenv("QDRANT_COLLECTION", "rag_knowledge")
|
||||
|
||||
EMBEDDING_MODEL: str = os.getenv("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text")
|
||||
TEXT_SPLITTER_PARAMS = {
|
||||
"chunk_size": 1000,
|
||||
"chunk_overlap": 200,
|
||||
}
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Класс клиента Qdrant
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class QdrantClient:
|
||||
"""
|
||||
Класс-обёртка над QdrantVectorStore.
|
||||
Обёртка над qdrant_client и langchain_qdrant для упрощения работы с коллекцией.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: str = "knowledge_base",
|
||||
host: str | None = None,
|
||||
port: int | None = None,
|
||||
embedding_model: str = "nomic-embed-text",
|
||||
chunk_size: int = 1000,
|
||||
chunk_overlap: int = 200,
|
||||
) -> None:
|
||||
"""
|
||||
Инициализация клиента Qdrant и подготовка коллекции.
|
||||
|
||||
:param collection_name: имя коллекции в Qdrant
|
||||
:param host: хост Qdrant (по умолчанию localhost)
|
||||
:param port: порт Qdrant (по умолчанию 6333)
|
||||
:param embedding_model: название модели Ollama для эмбеддингов
|
||||
:param chunk_size: максимальная длина чанка в символах
|
||||
:param chunk_overlap: перекрытие между чанками
|
||||
"""
|
||||
def __init__(self, host: str = QDRANT_HOST, port: int = QDRANT_PORT,
|
||||
collection_name: str = COLLECTION_NAME) -> None:
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.collection_name = collection_name
|
||||
|
||||
# Параметры подключения к Qdrant
|
||||
host = host or os.getenv("QDRANT_HOST", "localhost")
|
||||
port = port or int(os.getenv("QDRANT_PORT", 6333))
|
||||
# Инициализируем низкоуровневый клиент
|
||||
self.client = _QdrantClient(host=self.host, port=self.port)
|
||||
|
||||
# Создаём объект эмбеддинговой модели Ollama
|
||||
self.embeddings = OllamaEmbeddings(model=embedding_model)
|
||||
# Создаём коллекцию при необходимости
|
||||
if not self._collection_exists():
|
||||
self._create_collection()
|
||||
|
||||
# Инициализируем QdrantVectorStore
|
||||
self.store = QdrantVectorStore(
|
||||
url=f"http://{host}:{port}",
|
||||
# Обёртка LangChain для упрощённого добавления/поиска
|
||||
self.vector_store = QdrantVectorStore(
|
||||
client=self.client,
|
||||
collection_name=self.collection_name,
|
||||
embedding_function=self.embeddings.embed_query, # используем embed_query для совместимости
|
||||
embedding=OllamaEmbeddings(model=EMBEDDING_MODEL),
|
||||
)
|
||||
|
||||
# Создаём splitter для разбиения текста на чанки
|
||||
self.splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
# ----------------------------------------------------------------------- #
|
||||
# Внутренние вспомогательные методы
|
||||
# ----------------------------------------------------------------------- #
|
||||
|
||||
def _collection_exists(self) -> bool:
|
||||
"""Проверяем наличие коллекции."""
|
||||
try:
|
||||
self.client.get_collection(name=self.collection_name)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _create_collection(self) -> None:
|
||||
"""Создаём коллекцию с параметрами по умолчанию."""
|
||||
vector_params = VectorParams(size=384, distance=Distance.COSINE)
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=vector_params
|
||||
)
|
||||
|
||||
def add_document(self, content: str, title: str) -> None:
|
||||
"""
|
||||
Добавляет документ в базу знаний. Текст разбивается на чанки,
|
||||
каждому присваиваются метаданные (title и номер чанка).
|
||||
# ----------------------------------------------------------------------- #
|
||||
# Публичные методы
|
||||
# ----------------------------------------------------------------------- #
|
||||
|
||||
:param content: полный текст документа
|
||||
:param title: заголовок/имя документа
|
||||
def add_documents(self, documents: Iterable[str], titles: Iterable[str]) -> None:
|
||||
"""
|
||||
# Разбиваем контент на чанки
|
||||
chunks = self.splitter.split_text(content)
|
||||
Добавляет документы в коллекцию после разбиения на чанки.
|
||||
|
||||
documents = []
|
||||
for idx, chunk in enumerate(chunks):
|
||||
doc = {
|
||||
"page_content": chunk,
|
||||
"metadata": {"title": title, "chunk_index": idx},
|
||||
:param documents: Итерируемый список текстов.
|
||||
:param titles: Заголовки документов (один к одному с documents).
|
||||
"""
|
||||
splitter = RecursiveCharacterTextSplitter(**TEXT_SPLITTER_PARAMS)
|
||||
all_chunks: List[Tuple[str, str]] = []
|
||||
|
||||
for doc_text, title in zip(documents, titles):
|
||||
chunks = splitter.split_text(doc_text)
|
||||
# Для каждого чанка сохраняем заголовок как метаданные
|
||||
all_chunks.extend([(chunk, title) for chunk in chunks])
|
||||
|
||||
texts, metas = zip(*all_chunks)
|
||||
|
||||
self.vector_store.add_texts(
|
||||
texts=texts,
|
||||
metadatas=[{"title": meta} for meta in metas]
|
||||
)
|
||||
|
||||
def search(self, query: str, limit: int = 5) -> List[dict]:
|
||||
"""
|
||||
Семантический поиск по запросу.
|
||||
|
||||
:param query: Текст запроса.
|
||||
:param limit: Количество результатов.
|
||||
:return: Список словарей с полями 'text', 'score' и метаданными.
|
||||
"""
|
||||
results = self.vector_store.similarity_search_with_score(query, k=limit)
|
||||
return [
|
||||
{
|
||||
"text": r[0].page_content,
|
||||
"score": r[1],
|
||||
"metadata": r[0].metadata
|
||||
}
|
||||
documents.append(doc)
|
||||
for r in results
|
||||
]
|
||||
|
||||
# Добавляем векторные представления чанков
|
||||
self.store.add_documents(documents)
|
||||
# ----------------------------------------------------------------------- #
|
||||
# Утилиты для загрузки файлов из директории
|
||||
# ----------------------------------------------------------------------- #
|
||||
|
||||
def search(
|
||||
self,
|
||||
query: str,
|
||||
max_results: int = 5,
|
||||
filter_metadata: Dict[str, Any] | None = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
def load_from_directory(self, directory: str | Path) -> None:
|
||||
"""
|
||||
Семантический поиск по базе знаний.
|
||||
Загружает все текстовые файлы из указанной папки в коллекцию.
|
||||
|
||||
:param query: поисковый запрос
|
||||
:param max_results: максимальное количество результатов
|
||||
:param filter_metadata: словарь метаданных для фильтрации (например, {"title": "MyDoc"})
|
||||
:return: список найденных документов с полями:
|
||||
- page_content
|
||||
- metadata
|
||||
- score
|
||||
:param directory: Путь к каталогу.
|
||||
"""
|
||||
# Выполняем поиск векторного пространства
|
||||
results = self.store.similarity_search_with_score(
|
||||
query=query,
|
||||
k=max_results,
|
||||
filter=filter_metadata,
|
||||
)
|
||||
dir_path = Path(directory)
|
||||
if not dir_path.is_dir():
|
||||
raise ValueError(f"Папка {directory} не существует")
|
||||
|
||||
# Преобразуем к удобному формату
|
||||
formatted = []
|
||||
for doc, score in results:
|
||||
formatted.append(
|
||||
{
|
||||
"content": doc.page_content,
|
||||
"metadata": doc.metadata,
|
||||
"score": score,
|
||||
}
|
||||
)
|
||||
return formatted
|
||||
docs = []
|
||||
titles = []
|
||||
|
||||
for file in dir_path.glob("**/*.txt"):
|
||||
text = file.read_text(encoding="utf-8")
|
||||
docs.append(text)
|
||||
titles.append(file.stem)
|
||||
|
||||
self.add_documents(docs, titles)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Пример использования (не входит в публичный API модуля)
|
||||
# Пример использования (можно удалить при импорте в другие модули)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Инициализируем клиент
|
||||
client = QdrantClient()
|
||||
# Загрузка из папки ./data
|
||||
client.load_from_directory("./data")
|
||||
|
||||
# Добавляем пример документа из файла
|
||||
sample_path = Path("sample.txt")
|
||||
if sample_path.exists():
|
||||
text = sample_path.read_text(encoding="utf-8")
|
||||
client.add_document(content=text, title=sample_path.stem)
|
||||
|
||||
# Выполняем поиск
|
||||
query_str = "What is the main idea of the document?"
|
||||
results = client.search(query=query_str, max_results=3)
|
||||
# Поиск по запросу
|
||||
query_text = "Какой смысл жизни?"
|
||||
results = client.search(query_text, limit=3)
|
||||
for i, res in enumerate(results, 1):
|
||||
print(f"\nResult {i} (score={res['score']:.4f})")
|
||||
print(f"\nРезультат {i}:")
|
||||
print(f"Score: {res['score']:.4f}")
|
||||
print(f"Title: {res['metadata'].get('title')}")
|
||||
print(f"Chunk index: {res['metadata'].get('chunk_index')}")
|
||||
print("Content snippet:", res["content"][:200], "...")
|
||||
print(f"Text snippet: {res['text'][:200]}...")
|
||||
Reference in New Issue
Block a user