restore own solution: vector_store.py
This commit is contained in:
@@ -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)
|
||||||
Reference in New Issue
Block a user