add vector_store.py

vector_store.py

@@ -0,0 +1,77 @@
+"""Векторное хранилище Qdrant + эмбеддинги Ollama."""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from langchain_core.documents import Document
+from langchain_ollama import OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import Distance, VectorParams
+
+COLLECTION_NAME = "knowledge_base"
+QDRANT_PATH = os.getenv("QDRANT_PATH", "./qdrant_data")
+OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
+EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
+
+
+def get_embeddings() -> OllamaEmbeddings:
+    return OllamaEmbeddings(
+        model=EMBED_MODEL,
+        base_url=OLLAMA_BASE_URL,
+    )
+
+
+def get_qdrant_client() -> QdrantClient:
+    Path(QDRANT_PATH).mkdir(parents=True, exist_ok=True)
+    return QdrantClient(path=QDRANT_PATH)
+
+
+def _ensure_collection(client: QdrantClient, embeddings: OllamaEmbeddings) -> None:
+    if client.collection_exists(COLLECTION_NAME):
+        return
+    vector_size = len(embeddings.embed_query("dimension probe"))
+    client.create_collection(
+        collection_name=COLLECTION_NAME,
+        vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE),
+    )
+
+
+def get_vector_store() -> QdrantVectorStore:
+    client = get_qdrant_client()
+    embeddings = get_embeddings()
+    _ensure_collection(client, embeddings)
+    return QdrantVectorStore(
+        client=client,
+        collection_name=COLLECTION_NAME,
+        embedding=embeddings,
+    )
+
+
+def chunk_document(content: str, title: str) -> list[Document]:
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=500,
+        chunk_overlap=50,
+    )
+    documents = splitter.create_documents(
+        texts=[content],
+        metadatas=[{"title": title}],
+    )
+    for index, doc in enumerate(documents):
+        doc.metadata["chunk_index"] = index
+        doc.metadata["source"] = title
+    return documents
+
+
+def add_document_to_store(content: str, title: str) -> int:
+    store = get_vector_store()
+    chunks = chunk_document(content, title)
+    store.add_documents(chunks)
+    return len(chunks)
+
+
+def search_store(query: str, max_results: int = 4) -> list[tuple[Document, float]]:
+    store = get_vector_store()
+    return store.similarity_search_with_score(query, k=max_results)