rework: ChromaDB - update vector_store.py

This commit is contained in:
2026-05-27 09:38:11 +00:00
parent b2c437bf4d
commit f4d63adda8
+8 -27
View File
@@ -1,18 +1,16 @@
"""Векторное хранилище Qdrant + эмбеддинги Ollama.""" """Векторное хранилище ChromaDB + эмбеддинги Ollama."""
from __future__ import annotations from __future__ import annotations
import os import os
from pathlib import Path from pathlib import Path
from langchain_chroma import Chroma
from langchain_core.documents import Document from langchain_core.documents import Document
from langchain_ollama import OllamaEmbeddings from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_text_splitters import RecursiveCharacterTextSplitter
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
COLLECTION_NAME = "knowledge_base" COLLECTION_NAME = "knowledge_base"
QDRANT_PATH = os.getenv("QDRANT_PATH", "./qdrant_data") CHROMA_PATH = os.getenv("CHROMA_PATH", "./chroma_data")
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434") OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text") EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
@@ -24,29 +22,12 @@ def get_embeddings() -> OllamaEmbeddings:
) )
def get_qdrant_client() -> QdrantClient: def get_vector_store() -> Chroma:
Path(QDRANT_PATH).mkdir(parents=True, exist_ok=True) Path(CHROMA_PATH).mkdir(parents=True, exist_ok=True)
return QdrantClient(path=QDRANT_PATH) return Chroma(
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, collection_name=COLLECTION_NAME,
vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE), embedding_function=get_embeddings(),
) persist_directory=CHROMA_PATH,
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,
) )