Обновить vector_store.py

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"
+COLLECTION_NAME = "knowledge_base"
 EMBEDDING_MODEL = "nomic-embed-text"
-EMBEDDING_DIM     = 768          # размер вектора nomic-embed-text
-
-CHUNK_SIZE        = 500
-CHUNK_OVERLAP     = 100
+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
+    output = []
+    for doc, score in results:
+        output.append(
+            {
+                "content": doc.page_content,
+                "metadata": doc.metadata,
+                "score": round(score, 4),
+            }
+        )
+    return output