From 7825305fb0454ce83afd84a7330a29d3e087f6fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9C=D0=B0=D1=80=D0=B8=D1=8F=20=D0=91=D0=B5=D1=80=D0=B4?= =?UTF-8?q?=D0=BD=D0=B8=D0=BA=D0=BE=D0=B2=D0=B0?= Date: Thu, 28 May 2026 13:06:27 +0000 Subject: [PATCH] =?UTF-8?q?=D0=9E=D0=B1=D0=BD=D0=BE=D0=B2=D0=B8=D1=82?= =?UTF-8?q?=D1=8C=20vector=5Fstore.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- vector_store.py | 99 ++++++++++++++++++------------------------------- 1 file changed, 36 insertions(+), 63 deletions(-) diff --git a/vector_store.py b/vector_store.py index 185a9d8..ba0a8b3 100644 --- a/vector_store.py +++ b/vector_store.py @@ -1,100 +1,73 @@ -""" -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_ollama import OllamaEmbeddings from langchain_core.documents import Document -from langchain_text_splitters import RecursiveCharacterTextSplitter from qdrant_client import QdrantClient from qdrant_client.models import Distance, VectorParams +from langchain.text_splitter import RecursiveCharacterTextSplitter +import uuid -# ── Конфигурация ────────────────────────────────────────────────────────────── -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 +COLLECTION_NAME = "knowledge_base" +EMBEDDING_MODEL = "nomic-embed-text" +VECTOR_SIZE = 768 -# ── Эмбеддинги через Ollama ─────────────────────────────────────────────────── -embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) +def get_embeddings(): + return OllamaEmbeddings(model=EMBEDDING_MODEL) -# ── Qdrant-клиент и коллекция ───────────────────────────────────────────────── -def _get_client() -> QdrantClient: - return QdrantClient(url=QDRANT_URL) +def get_qdrant_client(): + return QdrantClient(host="localhost", port=6333) -def init_collection() -> None: - """Создаёт коллекцию в Qdrant, если её ещё нет.""" - client = _get_client() - existing = [c.name for c in client.get_collections().collections] - if COLLECTION_NAME not in existing: +def init_collection(client: QdrantClient): + collections = [c.name for c in client.get_collections().collections] + if COLLECTION_NAME not in collections: client.create_collection( collection_name=COLLECTION_NAME, - vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE), + vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE), ) - print(f"[vector_store] Коллекция «{COLLECTION_NAME}» создана.") - else: - print(f"[vector_store] Коллекция «{COLLECTION_NAME}» уже существует.") -def _get_store() -> QdrantVectorStore: - """Возвращает готовый QdrantVectorStore.""" +def get_vector_store() -> QdrantVectorStore: + client = get_qdrant_client() + init_collection(client) + embeddings = get_embeddings() return QdrantVectorStore( - client=_get_client(), + client=client, collection_name=COLLECTION_NAME, embedding=embeddings, ) -# ── Чанкинг ─────────────────────────────────────────────────────────────────── -def split_text(content: str, title: str = "") -> List[Document]: - """ - Разбивает текст на чанки с помощью RecursiveCharacterTextSplitter - и добавляет метаданные (title, chunk_index). - """ +def add_documents(content: str, title: str) -> int: splitter = RecursiveCharacterTextSplitter( - chunk_size=CHUNK_SIZE, - chunk_overlap=CHUNK_OVERLAP, + chunk_size=500, + chunk_overlap=50, separators=["\n\n", "\n", ".", " ", ""], ) chunks = splitter.split_text(content) docs = [ Document( page_content=chunk, - metadata={"title": title, "chunk_index": i, "source": title or "manual"}, + metadata={"title": title, "chunk_index": i, "source": title}, ) 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 = get_vector_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() +def search_documents(query: str, max_results: int = 5) -> list[dict]: + store = get_vector_store() results = store.similarity_search_with_relevance_scores(query, k=max_results) - return results \ No newline at end of file + output = [] + for doc, score in results: + output.append( + { + "content": doc.page_content, + "metadata": doc.metadata, + "score": round(score, 4), + } + ) + return output \ No newline at end of file