55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
"""
|
|
Vector store abstraction using ChromaDB.
|
|
|
|
Provides methods to add documents with embeddings and perform similarity search.
|
|
"""
|
|
import os
|
|
from pathlib import Path
|
|
from typing import List, Dict, Any
|
|
|
|
from chromadb import Client as ChromaClient
|
|
from chromadb.config import Settings
|
|
from langchain_ollama import OllamaEmbeddings
|
|
|
|
# Initialize global client (in-memory for simplicity)
|
|
client = ChromaClient(Settings(chroma_db_impl="duckdb+parquet", persist_directory=None))
|
|
collection_name = "rag_collection"
|
|
|
|
# Ensure collection exists
|
|
if collection_name not in client.list_collections():
|
|
client.create_collection(name=collection_name)
|
|
col = client.get_or_create_collection(name=collection_name)
|
|
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
|
|
class ChromaVectorStore:
|
|
"""Wrapper around a Chroma collection."""
|
|
|
|
def __init__(self, collection):
|
|
self.collection = collection
|
|
|
|
def add_documents(self, documents: List[Dict[str, Any]]):
|
|
ids = []
|
|
texts = []
|
|
metadatas = []
|
|
for doc in documents:
|
|
ids.append(doc.get("id", os.urandom(8).hex()))
|
|
texts.append(doc["content"])
|
|
metadatas.append(doc.get("metadata", {}))
|
|
embeddings_list = embeddings.embed_documents(texts)
|
|
self.collection.add(ids=ids, documents=texts, embeddings=embeddings_list, metadatas=metadatas)
|
|
|
|
def similarity_search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
|
|
results = self.collection.query(
|
|
query_texts=[query],
|
|
n_results=k,
|
|
include=['documents', 'distances', 'metadatas'],
|
|
)
|
|
docs = []
|
|
for doc, dist, meta in zip(results["documents"][0], results["distances"][0], results["metadatas"][0]):
|
|
docs.append({"content": doc, "distance": dist, "metadata": meta})
|
|
return docs
|
|
|
|
# Singleton instance
|
|
vector_store = ChromaVectorStore(col)
|