100 lines
4.1 KiB
Python
100 lines
4.1 KiB
Python
"""
|
|
vector_store.py — модуль для работы с векторным хранилищем Qdrant + Ollama.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from typing import List, Tuple
|
|
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_qdrant import QdrantVectorStore
|
|
from langchain_core.documents import Document
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import Distance, VectorParams
|
|
|
|
# ── Конфигурация ──────────────────────────────────────────────────────────────
|
|
QDRANT_URL = "http://localhost:6333"
|
|
COLLECTION_NAME = "rag_memory"
|
|
EMBEDDING_MODEL = "nomic-embed-text"
|
|
EMBEDDING_DIM = 768 # размер вектора nomic-embed-text
|
|
|
|
CHUNK_SIZE = 500
|
|
CHUNK_OVERLAP = 100
|
|
|
|
|
|
# ── Эмбеддинги через Ollama ───────────────────────────────────────────────────
|
|
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
|
|
|
|
|
# ── Qdrant-клиент и коллекция ─────────────────────────────────────────────────
|
|
def _get_client() -> QdrantClient:
|
|
return QdrantClient(url=QDRANT_URL)
|
|
|
|
|
|
def init_collection() -> None:
|
|
"""Создаёт коллекцию в Qdrant, если её ещё нет."""
|
|
client = _get_client()
|
|
existing = [c.name for c in client.get_collections().collections]
|
|
if COLLECTION_NAME not in existing:
|
|
client.create_collection(
|
|
collection_name=COLLECTION_NAME,
|
|
vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE),
|
|
)
|
|
print(f"[vector_store] Коллекция «{COLLECTION_NAME}» создана.")
|
|
else:
|
|
print(f"[vector_store] Коллекция «{COLLECTION_NAME}» уже существует.")
|
|
|
|
|
|
def _get_store() -> QdrantVectorStore:
|
|
"""Возвращает готовый QdrantVectorStore."""
|
|
return QdrantVectorStore(
|
|
client=_get_client(),
|
|
collection_name=COLLECTION_NAME,
|
|
embedding=embeddings,
|
|
)
|
|
|
|
|
|
# ── Чанкинг ───────────────────────────────────────────────────────────────────
|
|
def split_text(content: str, title: str = "") -> List[Document]:
|
|
"""
|
|
Разбивает текст на чанки с помощью RecursiveCharacterTextSplitter
|
|
и добавляет метаданные (title, chunk_index).
|
|
"""
|
|
splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=CHUNK_SIZE,
|
|
chunk_overlap=CHUNK_OVERLAP,
|
|
separators=["\n\n", "\n", ".", " ", ""],
|
|
)
|
|
chunks = splitter.split_text(content)
|
|
docs = [
|
|
Document(
|
|
page_content=chunk,
|
|
metadata={"title": title, "chunk_index": i, "source": title or "manual"},
|
|
)
|
|
for i, chunk in enumerate(chunks)
|
|
]
|
|
return docs
|
|
|
|
|
|
# ── Публичный API ─────────────────────────────────────────────────────────────
|
|
def add_documents(content: str, title: str = "") -> int:
|
|
"""
|
|
Добавляет документ в векторное хранилище.
|
|
Возвращает количество добавленных чанков.
|
|
"""
|
|
docs = split_text(content, title)
|
|
store = _get_store()
|
|
store.add_documents(docs)
|
|
return len(docs)
|
|
|
|
|
|
def search(query: str, max_results: int = 5) -> List[Tuple[Document, float]]:
|
|
"""
|
|
Семантический поиск с метрикой релевантности.
|
|
Возвращает список (Document, score).
|
|
"""
|
|
store = _get_store()
|
|
results = store.similarity_search_with_relevance_scores(query, k=max_results)
|
|
return results |