diff --git a/vector_store.py b/vector_store.py new file mode 100644 index 0000000..5e734d8 --- /dev/null +++ b/vector_store.py @@ -0,0 +1,65 @@ +"""Векторное хранилище Qdrant + эмбеддинги Ollama.""" +from __future__ import annotations + +import os + +from langchain_qdrant import QdrantVectorStore +from langchain_core.documents import Document +from langchain_ollama import OllamaEmbeddings +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_vector_store() -> QdrantVectorStore: + client = QdrantClient(path=QDRANT_PATH) + collections = client.get_collections().collections + if not any(col.name == COLLECTION_NAME for col in collections): + client.create_collection( + collection_name=COLLECTION_NAME, + vectors_config=VectorParams(size=768, distance=Distance.COSINE), + ) + return QdrantVectorStore( + client=client, + collection_name=COLLECTION_NAME, + embedding=get_embeddings(), + ) + + +def chunk_document(content: str, title: str) -> list[Document]: + splitter = RecursiveCharacterTextSplitter( + chunk_size=500, + chunk_overlap=100, + ) + 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)