"""
Abstract vector store interface with Chroma backend.
chromadb is a core dependency. Chroma operates in two modes:
- Embedded PersistentClient: no server needed, local persistence
- HttpClient: connects to the Chroma Docker service (docker compose up -d)
Follows the D-014 pattern (abstract base + factory function).
Decision: D-017
"""
from __future__ import annotations
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
LOG = logging.getLogger("rag.vector_store")
[docs]
@dataclass
class VectorDocument:
"""A document stored in the vector store."""
id: str
text: str
embedding: list[float] | None = None
metadata: dict[str, Any] = field(default_factory=dict)
[docs]
@dataclass
class VectorSearchResult:
"""A single search result from the vector store."""
id: str
text: str
score: float
metadata: dict[str, Any] = field(default_factory=dict)
[docs]
class VectorStore(ABC):
"""
Abstract interface for vector storage and similarity search.
Implementations persist document embeddings and support
approximate nearest-neighbor queries.
"""
[docs]
@abstractmethod
def add(
self,
documents: list[VectorDocument],
) -> None:
"""Add documents to the store. Upserts on matching IDs."""
[docs]
@abstractmethod
def search(
self,
query_embedding: list[float],
top_k: int = 10,
filter_metadata: dict[str, Any] | None = None,
) -> list[VectorSearchResult]:
"""Search for similar documents by embedding vector."""
[docs]
@abstractmethod
def delete(self, ids: list[str]) -> None:
"""Delete documents by ID."""
[docs]
@abstractmethod
def count(self) -> int:
"""Return the number of documents in the store."""
[docs]
def close(self) -> None:
"""Release resources. Override if needed."""
pass
[docs]
class ChromaVectorStore(VectorStore):
"""
Chroma-based vector store (core dependency).
Operates in two modes:
- Embedded (PersistentClient): no server, local persistence
- Client/server (HttpClient): connects to docker compose chroma service
Set CHROMA_HOST env var to connect to a remote Chroma instance,
otherwise defaults to embedded mode.
Decision: D-017
"""
def __init__(
self,
collection_name: str = "code_nodes",
persist_directory: str | None = None,
chroma_host: str | None = None,
chroma_port: int = 8000,
) -> None:
import chromadb
if chroma_host:
self._client = chromadb.HttpClient(host=chroma_host, port=chroma_port)
LOG.info("Chroma: connected to %s:%d", chroma_host, chroma_port)
elif persist_directory:
Path(persist_directory).mkdir(parents=True, exist_ok=True)
self._client = chromadb.PersistentClient(path=persist_directory)
LOG.info("Chroma: persistent at %s", persist_directory)
else:
self._client = chromadb.Client()
LOG.info("Chroma: ephemeral (in-memory)")
self._collection = self._client.get_or_create_collection(
name=collection_name,
metadata={"hnsw:space": "cosine"},
)
[docs]
def add(self, documents: list[VectorDocument]) -> None:
if not documents:
return
ids = [d.id for d in documents]
texts = [d.text for d in documents]
metadatas = [d.metadata if d.metadata else {"_placeholder": "true"} for d in documents]
kwargs: dict[str, Any] = {
"ids": ids,
"documents": texts,
"metadatas": metadatas,
}
# Include embeddings if provided
embeddings = [d.embedding for d in documents if d.embedding is not None]
if len(embeddings) == len(documents):
kwargs["embeddings"] = embeddings
self._collection.upsert(**kwargs)
[docs]
def search(
self,
query_embedding: list[float],
top_k: int = 10,
filter_metadata: dict[str, Any] | None = None,
) -> list[VectorSearchResult]:
n = min(top_k, self._collection.count())
if n == 0:
return []
kwargs: dict[str, Any] = {
"query_embeddings": [query_embedding],
"n_results": n,
}
if filter_metadata:
kwargs["where"] = filter_metadata
results = self._collection.query(**kwargs)
out: list[VectorSearchResult] = []
if results and results["ids"]:
for i, doc_id in enumerate(results["ids"][0]):
distance = results["distances"][0][i] if results.get("distances") else 0.0
score = max(0.0, 1.0 - distance) # cosine distance → similarity
text = results["documents"][0][i] if results.get("documents") else ""
meta = results["metadatas"][0][i] if results.get("metadatas") else {}
out.append(VectorSearchResult(id=doc_id, text=text, score=score, metadata=meta))
return out
[docs]
def delete(self, ids: list[str]) -> None:
if ids:
self._collection.delete(ids=ids)
[docs]
def count(self) -> int:
return self._collection.count()
[docs]
def build_vector_store(
backend: str = "chroma",
**kwargs: Any,
) -> VectorStore:
"""
Factory: create a VectorStore of the requested type.
Args:
backend: "chroma" (only supported backend currently)
**kwargs: Backend-specific configuration
Returns:
VectorStore instance
Raises:
ValueError: Unknown backend
"""
if backend == "chroma":
return ChromaVectorStore(**kwargs)
else:
raise ValueError(f"Unknown vector store backend: {backend!r}. Supported: 'chroma'")