62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
"""ChromaDB vector store with Ollama embeddings and document chunking.""
|
|
import os
|
|
from typing import List
|
|
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_core.documents import Document
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
|
|
COLLECTION_NAME = "knowledge_base"
|
|
PERSIST_DIR = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db")
|
|
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
|
|
_store: Chroma | None = None
|
|
|
|
|
|
def get_vector_store() -> Chroma:
|
|
"""Return singleton ChromaDB instance (persisted to disk)."""
|
|
global _store
|
|
if _store is None:
|
|
_store = Chroma(
|
|
collection_name=COLLECTION_NAME,
|
|
embedding_function=embeddings,
|
|
persist_directory=PERSIST_DIR,
|
|
)
|
|
return _store
|
|
|
|
|
|
def add_documents(texts: List[str], chunk_size: int = 500, chunk_overlap: int = 100) -> int:
|
|
"""Split texts into chunks and index them in ChromaDB.
|
|
|
|
Args:
|
|
texts: list of raw text strings
|
|
chunk_size: maximum chunk length in characters
|
|
chunk_overlap: overlap between consecutive chunks
|
|
|
|
Returns:
|
|
number of chunks added to the store
|
|
"""
|
|
splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=chunk_size,
|
|
chunk_overlap=chunk_overlap,
|
|
)
|
|
docs: List[Document] = splitter.create_documents(texts)
|
|
if docs:
|
|
get_vector_store().add_documents(docs)
|
|
return len(docs)
|
|
|
|
|
|
def similarity_search(query: str, k: int = 5) -> List[Document]:
|
|
"""Semantic search in ChromaDB.
|
|
|
|
Args:
|
|
query: natural language search query
|
|
k: number of top results to return
|
|
|
|
Returns:
|
|
list of the most relevant Document objects
|
|
"""
|
|
return get_vector_store().similarity_search(query, k=k)
|