rework: ChromaDB - update vector_store.py
This commit is contained in:
+8
-27
@@ -1,18 +1,16 @@
|
||||
"""Векторное хранилище Qdrant + эмбеддинги Ollama."""
|
||||
"""Векторное хранилище ChromaDB + эмбеддинги Ollama."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from langchain_chroma import Chroma
|
||||
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")
|
||||
CHROMA_PATH = os.getenv("CHROMA_PATH", "./chroma_data")
|
||||
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
|
||||
EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
|
||||
|
||||
@@ -24,29 +22,12 @@ def get_embeddings() -> OllamaEmbeddings:
|
||||
)
|
||||
|
||||
|
||||
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(
|
||||
def get_vector_store() -> Chroma:
|
||||
Path(CHROMA_PATH).mkdir(parents=True, exist_ok=True)
|
||||
return Chroma(
|
||||
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,
|
||||
embedding_function=get_embeddings(),
|
||||
persist_directory=CHROMA_PATH,
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user