175 lines
7.0 KiB
Python
175 lines
7.0 KiB
Python
# qdrant_client.py
|
|
"""
|
|
Модуль инициализации клиента Qdrant и простых операций с ним.
|
|
Используется в агенте для RAG‑памяти.
|
|
|
|
Поддержка:
|
|
- Создание/получение коллекции
|
|
- Добавление документов (с чанками) с эмбеддингами Ollama
|
|
- Семантический поиск по запросу
|
|
|
|
Требуемые зависимости:
|
|
pip install langchain langchain-qdrant langchain-ollama qdrant-client
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
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:
|
|
"""
|
|
Обёртка над qdrant_client и langchain_qdrant для упрощения работы с коллекцией.
|
|
"""
|
|
|
|
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
|
|
|
|
# Инициализируем низкоуровневый клиент
|
|
self.client = _QdrantClient(host=self.host, port=self.port)
|
|
|
|
# Создаём коллекцию при необходимости
|
|
if not self._collection_exists():
|
|
self._create_collection()
|
|
|
|
# Обёртка LangChain для упрощённого добавления/поиска
|
|
self.vector_store = QdrantVectorStore(
|
|
client=self.client,
|
|
collection_name=self.collection_name,
|
|
embedding=OllamaEmbeddings(model=EMBEDDING_MODEL),
|
|
)
|
|
|
|
# ----------------------------------------------------------------------- #
|
|
# Внутренние вспомогательные методы
|
|
# ----------------------------------------------------------------------- #
|
|
|
|
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_documents(self, documents: Iterable[str], titles: Iterable[str]) -> None:
|
|
"""
|
|
Добавляет документы в коллекцию после разбиения на чанки.
|
|
|
|
: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
|
|
}
|
|
for r in results
|
|
]
|
|
|
|
# ----------------------------------------------------------------------- #
|
|
# Утилиты для загрузки файлов из директории
|
|
# ----------------------------------------------------------------------- #
|
|
|
|
def load_from_directory(self, directory: str | Path) -> None:
|
|
"""
|
|
Загружает все текстовые файлы из указанной папки в коллекцию.
|
|
|
|
:param directory: Путь к каталогу.
|
|
"""
|
|
dir_path = Path(directory)
|
|
if not dir_path.is_dir():
|
|
raise ValueError(f"Папка {directory} не существует")
|
|
|
|
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)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Пример использования (можно удалить при импорте в другие модули)
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
if __name__ == "__main__":
|
|
client = QdrantClient()
|
|
# Загрузка из папки ./data
|
|
client.load_from_directory("./data")
|
|
|
|
# Поиск по запросу
|
|
query_text = "Какой смысл жизни?"
|
|
results = client.search(query_text, limit=3)
|
|
for i, res in enumerate(results, 1):
|
|
print(f"\nРезультат {i}:")
|
|
print(f"Score: {res['score']:.4f}")
|
|
print(f"Title: {res['metadata'].get('title')}")
|
|
print(f"Text snippet: {res['text'][:200]}...") |