""" Parquet-backed embedding store, adapted from the original EmbeddingStore. Uses QAFD-RAG's async embedding functions (wrapped synchronously) so we can share models / GPU memory with the rest of the QAFD-RAG system. """ import asyncio import logging import os from copy import deepcopy from typing import List, Dict, Optional, Callable, Any import numpy as np import pandas as pd from .utils import compute_mdhash_id logger = logging.getLogger(__name__) class EmbeddingModelWrapper: """Thin sync wrapper around a QAFD-RAG *async* embedding function. The wrapped function must have the signature:: async def embed(texts: list[str], **kwargs) -> np.ndarray Parameters ---------- embed_func : callable An async embedding function from ``QAFD-RAG/src/llm.py``. batch_size : int Max texts per call. """ def __init__(self, embed_func: Callable, batch_size: int = 16): self._embed_func = embed_func self.batch_size = batch_size # ------------------------------------------------------------------ def batch_encode(self, texts, instruction: str = None, norm: bool = True) -> np.ndarray: """Synchronously encode *texts* into embeddings.""" if isinstance(texts, str): texts = [texts] all_embeddings = [] # Use larger batch for API-based embeddings (OpenAI supports up to 2048) effective_batch = max(self.batch_size, 512) for start in range(0, len(texts), effective_batch): batch = texts[start : start + effective_batch] if instruction: batch = [f"{instruction} {t}" for t in batch] embs = self._run_async(self._embed_func(batch)) if not isinstance(embs, np.ndarray): embs = np.array(embs) if norm: norms = np.linalg.norm(embs, axis=1, keepdims=True) norms = np.where(norms == 0, 1, norms) embs = embs / norms all_embeddings.append(embs) return np.vstack(all_embeddings) # ------------------------------------------------------------------ @staticmethod def _run_async(coro): """Run an async coroutine synchronously.""" try: loop = asyncio.get_running_loop() except RuntimeError: loop = None if loop is not None and loop.is_running(): # We are inside an already-running event loop (e.g. Jupyter). import concurrent.futures with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: return pool.submit(asyncio.run, coro).result() else: return asyncio.run(coro) class EmbeddingStore: """Parquet-backed vector store. Mirrors the original EmbeddingStore but uses ``EmbeddingModelWrapper`` (which calls QAFD-RAG's async embedding functions under the hood). """ def __init__( self, embedding_model: EmbeddingModelWrapper, db_filename: str, batch_size: int, namespace: str, ): self.embedding_model = embedding_model self.batch_size = batch_size self.namespace = namespace if not os.path.exists(db_filename): logger.info(f"Creating directory: {db_filename}") os.makedirs(db_filename, exist_ok=True) self.filename = os.path.join(db_filename, f"vdb_{self.namespace}.parquet") self._load_data() # ------------------------------------------------------------------ # Data persistence # ------------------------------------------------------------------ def _load_data(self): if os.path.exists(self.filename): df = pd.read_parquet(self.filename) self.hash_ids = df["hash_id"].values.tolist() self.texts = df["content"].values.tolist() self.embeddings = df["embedding"].values.tolist() self._rebuild_indices() assert len(self.hash_ids) == len(self.texts) == len(self.embeddings) logger.info(f"Loaded {len(self.hash_ids)} records from {self.filename}") else: self.hash_ids, self.texts, self.embeddings = [], [], [] self.hash_id_to_idx: Dict[str, int] = {} self.hash_id_to_row: Dict[str, dict] = {} self.hash_id_to_text: Dict[str, str] = {} self.text_to_hash_id: Dict[str, str] = {} def _rebuild_indices(self): self.hash_id_to_idx = {h: idx for idx, h in enumerate(self.hash_ids)} self.hash_id_to_row = { h: {"hash_id": h, "content": t} for h, t in zip(self.hash_ids, self.texts) } self.hash_id_to_text = {h: self.texts[idx] for idx, h in enumerate(self.hash_ids)} self.text_to_hash_id = {self.texts[idx]: h for idx, h in enumerate(self.hash_ids)} def _save_data(self): data = pd.DataFrame({ "hash_id": self.hash_ids, "content": self.texts, "embedding": self.embeddings, }) data.to_parquet(self.filename, index=False) self._rebuild_indices() logger.info(f"Saved {len(self.hash_ids)} records to {self.filename}") def _upsert(self, hash_ids, texts, embeddings): self.embeddings.extend(embeddings) self.hash_ids.extend(hash_ids) self.texts.extend(texts) self._save_data() # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ def get_missing_string_hash_ids(self, texts: List[str]) -> Dict[str, dict]: nodes_dict = {} for text in texts: hid = compute_mdhash_id(text, prefix=self.namespace + "-") nodes_dict[hid] = {"content": text} if not nodes_dict: return {} existing = set(self.hash_id_to_row.keys()) missing = {h: {"hash_id": h, "content": v["content"]} for h, v in nodes_dict.items() if h not in existing} return missing def insert_strings(self, texts: List[str]): nodes_dict = {} for text in texts: if not text or not text.strip(): continue hid = compute_mdhash_id(text, prefix=self.namespace + "-") nodes_dict[hid] = {"content": text} all_ids = list(nodes_dict.keys()) if not all_ids: return existing = set(self.hash_id_to_row.keys()) missing_ids = [h for h in all_ids if h not in existing] logger.info( f"Inserting {len(missing_ids)} new records, " f"{len(all_ids) - len(missing_ids)} already exist." ) if not missing_ids: return texts_to_encode = [nodes_dict[h]["content"] for h in missing_ids] missing_embeddings = self.embedding_model.batch_encode(texts_to_encode) # Convert ndarray rows to list of lists for parquet storage if isinstance(missing_embeddings, np.ndarray): missing_embeddings = missing_embeddings.tolist() self._upsert(missing_ids, texts_to_encode, missing_embeddings) def delete(self, hash_ids): indices = sorted( [self.hash_id_to_idx[h] for h in hash_ids], reverse=True ) for idx in indices: self.hash_ids.pop(idx) self.texts.pop(idx) self.embeddings.pop(idx) self._save_data() # Lookups def get_row(self, hash_id: str) -> dict: return self.hash_id_to_row[hash_id] def get_hash_id(self, text: str) -> str: return self.text_to_hash_id[text] def get_rows(self, hash_ids: List[str], dtype=np.float32) -> Dict[str, dict]: if not hash_ids: return {} return {hid: self.hash_id_to_row[hid] for hid in hash_ids} def get_all_ids(self) -> List[str]: return deepcopy(self.hash_ids) def get_all_id_to_rows(self) -> Dict[str, dict]: return deepcopy(self.hash_id_to_row) def get_all_texts(self) -> set: return set(row["content"] for row in self.hash_id_to_row.values()) def get_embedding(self, hash_id: str, dtype=np.float32) -> np.ndarray: return np.array(self.embeddings[self.hash_id_to_idx[hash_id]], dtype=dtype) def get_embeddings(self, hash_ids: List[str], dtype=np.float32) -> np.ndarray: if not hash_ids: return np.array([]) indices = np.array([self.hash_id_to_idx[h] for h in hash_ids], dtype=np.intp) all_embs = np.array(self.embeddings, dtype=dtype) return all_embs[indices]