diff --git a/vector_store.py b/vector_store.py index 5e734d8..be633fa 100644 --- a/vector_store.py +++ b/vector_store.py @@ -1,17 +1,16 @@ -"""Векторное хранилище Qdrant + эмбеддинги Ollama.""" +"""Векторное хранилище ChromaDB + эмбеддинги Ollama.""" from __future__ import annotations import os +from pathlib import Path -from langchain_qdrant import QdrantVectorStore +from langchain_chroma import Chroma 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") +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") @@ -23,18 +22,12 @@ def get_embeddings() -> OllamaEmbeddings: ) -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, +def get_vector_store() -> Chroma: + Path(CHROMA_PATH).mkdir(parents=True, exist_ok=True) + return Chroma( collection_name=COLLECTION_NAME, - embedding=get_embeddings(), + embedding_function=get_embeddings(), + persist_directory=CHROMA_PATH, )