Spaces:
Sleeping
Sleeping
File size: 4,454 Bytes
ce11d27 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """
ChromaDB vector store for persistent embedding storage.
Collections are named {session_id}__{model_name} so the same conversation
can be embedded with different models and compared.
"""
from __future__ import annotations
import logging
import threading
from typing import List, Optional
import numpy as np
import chromadb
logger = logging.getLogger(__name__)
# ChromaDB's Rust backend is not thread-safe for concurrent initialization.
# Serialize all PersistentClient creation to avoid segfaults / AttributeErrors.
_chromadb_init_lock = threading.Lock()
class VectorStore:
"""ChromaDB wrapper for storing and retrieving embeddings.
Args:
persist_dir: Directory for ChromaDB persistent storage.
"""
def __init__(self, persist_dir: str):
with _chromadb_init_lock:
self._client = chromadb.PersistentClient(path=persist_dir)
@staticmethod
def _collection_name(session_id: str, model_name: str) -> str:
"""Build a collection name from session ID and model name.
ChromaDB collection names must be 3-63 chars, start/end with
alphanumeric, and contain only alphanumerics, underscores, hyphens.
"""
raw = f"{session_id}__{model_name}"
sanitized = "".join(c if c.isalnum() or c in ("_", "-") else "_" for c in raw)
if len(sanitized) < 3:
sanitized = sanitized + "___"
return sanitized[:63]
def store_embeddings(
self,
session_id: str,
model_name: str,
texts: List[str],
embeddings: np.ndarray,
metadatas: Optional[List[dict]] = None,
):
"""Store embeddings for a session+model pair.
Args:
session_id: Session identifier.
model_name: Embedding model name.
texts: List of text strings (N).
embeddings: (N, D) array of embedding vectors.
metadatas: Optional per-entry metadata dicts.
"""
col_name = self._collection_name(session_id, model_name)
collection = self._client.get_or_create_collection(
name=col_name,
metadata={"session_id": session_id, "model_name": model_name},
)
ids = [f"{session_id}_{i}" for i in range(len(texts))]
if metadatas is None:
metadatas = [{"index": i} for i in range(len(texts))]
collection.upsert(
ids=ids,
documents=texts,
embeddings=embeddings.tolist(),
metadatas=metadatas,
)
def load_embeddings(
self, session_id: str, model_name: str
) -> Optional[np.ndarray]:
"""Load stored embeddings for a session+model pair.
Returns (N, D) array or None if not found.
"""
col_name = self._collection_name(session_id, model_name)
try:
collection = self._client.get_collection(name=col_name)
except Exception as e:
logger.debug(f"Collection not found: {col_name}: {e}")
return None
result = collection.get(include=["embeddings"])
if result["embeddings"] is None or len(result["embeddings"]) == 0:
return None
return np.array(result["embeddings"], dtype=np.float32)
def list_sessions(self) -> List[dict]:
"""List all stored session/model combinations."""
collections = self._client.list_collections()
sessions = []
for col in collections:
meta = col.metadata or {}
sessions.append({
"collection_name": col.name,
"session_id": meta.get("session_id", "unknown"),
"model_name": meta.get("model_name", "unknown"),
"count": col.count(),
})
return sessions
def delete_session(self, session_id: str, model_name: str):
"""Delete a stored session+model collection."""
col_name = self._collection_name(session_id, model_name)
try:
self._client.delete_collection(name=col_name)
except Exception as e:
logger.debug(f"Failed to delete collection {col_name}: {e}")
def clear_all(self):
"""Delete all stored embedding collections."""
for col in self._client.list_collections():
try:
self._client.delete_collection(name=col.name)
except Exception as e:
logger.debug(f"Failed to delete collection {col.name}: {e}")
|