66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
"""
|
||
Vector store implementation using ChromaDB.
|
||
|
||
Provides functions to add documents and perform semantic search.
|
||
"""
|
||
|
||
import os
|
||
from typing import List, Dict
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from chromadb import Client
|
||
from chromadb.config import Settings
|
||
|
||
# Initialize embeddings model (Ollama)
|
||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||
|
||
# ChromaDB client – in‑memory by default, persistent folder "chromadb"
|
||
CHROMA_DIR = os.path.join(os.getcwd(), "chromadb")
|
||
client = Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory=CHROMA_DIR))
|
||
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)
|
||
|
||
class VectorStore:
|
||
def __init__(self, collection):
|
||
self.collection = collection
|
||
|
||
def add_document(self, doc_id: str, text: str, metadata: Dict | None = None) -> None:
|
||
"""Add a single document to the collection.
|
||
|
||
Parameters
|
||
----------
|
||
doc_id: str
|
||
Unique identifier for the document.
|
||
text: str
|
||
Raw text content.
|
||
metadata: dict, optional
|
||
Additional key/value pairs stored with the vector.
|
||
"""
|
||
vec = embeddings.embed_query(text)
|
||
self.collection.add(ids=[doc_id], documents=[text], metadatas=[metadata or {}])
|
||
|
||
def search(self, query: str, k: int = 5) -> List[Dict]:
|
||
"""Semantic search over the collection.
|
||
|
||
Returns a list of dicts with keys: id, document, score, metadata.
|
||
"""
|
||
results = self.collection.query(
|
||
query_texts=[query], n_results=k, include=['documents', 'distances', 'metadatas']
|
||
)
|
||
hits = []
|
||
for i in range(len(results["ids"][0])):
|
||
hit = {
|
||
"id": results["ids"][0][i],
|
||
"document": results["documents"][0][i],
|
||
"score": 1 - results["distances"][0][i], # distance to similarity
|
||
"metadata": results["metadatas"][0][i],
|
||
}
|
||
hits.append(hit)
|
||
return hits
|
||
|
||
# Singleton instance for easy import
|
||
vector_store = VectorStore(col)
|