#!/usr/bin/env python3 """ TIER-1 VECTOR STORE - Supabase pgvector ======================================= Replaces brute-force Redis cosine search with proper ANN indexing. Uses Supabase pgvector extension - IVFFlat/HNSW indexes. Scales to millions of documents with sub-20ms queries. Architecture: - pgvector table with IVFFlat index (fast approximate search) - Hybrid search: dense (vector) + sparse (BM25 text) - Metadata filtering (chain, severity, date range) - Automatic index maintenance """ import json import logging import os from datetime import UTC, datetime from typing import Any import httpx from dotenv import load_dotenv # Load env vars at module level so keys are available before any class instantiation load_dotenv("/app/.env", override=True) logger = logging.getLogger(__name__) def _get_url(): """Get Supabase URL from env dynamically.""" return os.environ.get("SUPABASE_URL", "") def _get_headers(): """Build headers dynamically to pick up env changes without module reload.""" key = os.environ.get("SUPABASE_SERVICE_KEY", "") or os.environ.get("SUPABASE_SERVICE_ROLE_KEY", "") return { "apikey": key, "Authorization": f"Bearer {key}", "Content-Type": "application/json", } # Table configuration VECTOR_TABLE = "rag_vectors" # Dynamic embedding dimension - determined by the active model. # Local BGE-small = 384, BGE-M3/OpenRouter = 3072, etc. # Set via env var or detect at runtime from the embedder. EMBEDDING_DIM = int(os.environ.get("RAG_EMBEDDING_DIM", "0")) # 0 = auto-detect # The target dimension for the pgvector table column and RPC calls. # This MUST match the vector(N) in the CREATE TABLE and search_embeddings RPC. # Migration scripts use 640, .env sets RAG_EMBEDDING_DIM=640. # Defaults to EMBEDDING_DIM if set, otherwise 640. TABLE_DIM = int(os.environ.get("RAG_TABLE_DIM", str(max(EMBEDDING_DIM, 1)))) if EMBEDDING_DIM > 0 else 640 def pad_vector(vec: list, target_dim: int) -> list: """Pad or truncate a vector to exactly target_dim dimensions.""" if len(vec) == target_dim: return vec if len(vec) > target_dim: return vec[:target_dim] return vec + [0.0] * (target_dim - len(vec)) # SQL for table creation CREATE_TABLE_SQL = f""" CREATE TABLE IF NOT EXISTS {VECTOR_TABLE} ( id TEXT PRIMARY KEY, collection TEXT NOT NULL, content TEXT, embedding vector({TABLE_DIM}), metadata JSONB DEFAULT '{{}}', source TEXT, severity TEXT, chain TEXT, created_at TIMESTAMPTZ DEFAULT NOW(), updated_at TIMESTAMPTZ DEFAULT NOW() ); -- Indexes CREATE INDEX IF NOT EXISTS idx_rag_collection ON {VECTOR_TABLE} (collection); CREATE INDEX IF NOT EXISTS idx_rag_metadata ON {VECTOR_TABLE} USING GIN (metadata); CREATE INDEX IF NOT EXISTS idx_rag_severity ON {VECTOR_TABLE} (severity); CREATE INDEX IF NOT EXISTS idx_rag_chain ON {VECTOR_TABLE} (chain); CREATE INDEX IF NOT EXISTS idx_rag_source ON {VECTOR_TABLE} (source); -- Full-text search index for BM25-style keyword matching CREATE INDEX IF NOT EXISTS idx_rag_content_fts ON {VECTOR_TABLE} USING GIN (to_tsvector('english', COALESCE(content, ''))); -- IVFFlat index for approximate nearest neighbor search -- Note: Must be created after data is loaded for best results -- CREATE INDEX IF NOT EXISTS idx_rag_embedding_ivfflat ON {VECTOR_TABLE} -- USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100); """ class SupabaseVectorStore: """ Tier-1 vector store using Supabase pgvector. Supports: - ANN search with IVFFlat/HNSW - Hybrid search (dense + sparse) - Metadata filtering - Batch ingestion - Collection management """ def __init__(self): self._table_ready = False self._resolved_dim = EMBEDDING_DIM # may be 0 = auto-detect def _get_dim(self, embedding: list[float] | None = None) -> int: """Resolve the embedding dimension: env var > actual vector length > fallback 384.""" if self._resolved_dim > 0: return self._resolved_dim # Auto-detect from the first embedding we see if embedding and len(embedding) > 0: self._resolved_dim = len(embedding) logger.info(f"Auto-detected embedding dimension: {self._resolved_dim}") return self._resolved_dim # Default: local BGE-small-en-v1.5 = 384 self._resolved_dim = 384 logger.info(f"Using default embedding dimension: {self._resolved_dim}") return self._resolved_dim async def _rpc(self, fn: str, params: dict | None = None) -> dict: """Execute a Supabase RPC function.""" url = f"{_get_url()}/rest/v1/rpc/{fn}" async with httpx.AsyncClient(timeout=30) as client: resp = await client.post(url, json=params or {}, headers=_get_headers()) resp.raise_for_status() return resp.json() if resp.text else {} async def _query( self, table: str, select: str = "*", filters: dict | None = None, limit: int = 100, offset: int = 0, ) -> list[dict]: """Query Supabase REST API.""" url = f"{_get_url()}/rest/v1/{table}" params = {"select": select, "limit": str(limit), "offset": str(offset)} if filters: for k, v in filters.items(): params[k] = f"eq.{v}" async with httpx.AsyncClient(timeout=30) as client: resp = await client.get(url, params=params, headers=_get_headers()) resp.raise_for_status() return resp.json() if resp.text else [] async def _upsert(self, table: str, rows: list[dict]) -> dict: """Upsert rows into a Supabase table.""" url = f"{_get_url()}/rest/v1/{table}" params = {"on_conflict": "id"} async with httpx.AsyncClient(timeout=30) as client: resp = await client.post(url, json=rows, params=params, headers=_get_headers()) resp.raise_for_status() return resp.json() if resp.text else {} async def initialize(self) -> bool: """Create table and indexes if not exist. Falls back gracefully.""" global TABLE_DIM # Determine the query dimension for the RPC health check. # If the RPC function already exists on Supabase, it has a fixed signature # (e.g. vector(640) or vector(1024)). We probe with TABLE_DIM first, # then fall back to common alternatives. rpc_probed = False for probe_dim in [TABLE_DIM, 640, 1024, 384]: try: await self._rpc( "search_embeddings", { "query_embedding": [0.1] * probe_dim, "match_count": 1, "similarity_threshold": 0.0, }, ) self._table_ready = True TABLE_DIM = probe_dim logger.info(f"pgvector: search_embeddings RPC available (dim={probe_dim})") rpc_probed = True break except Exception: continue if not rpc_probed: logger.warning("pgvector: search_embeddings RPC not available") # Try to create table via SQL RPC if available if not self._table_ready: try: # Build the CREATE TABLE SQL with the resolved dimension dim = self._get_dim() create_sql = f""" CREATE TABLE IF NOT EXISTS {VECTOR_TABLE} ( id TEXT PRIMARY KEY, collection TEXT NOT NULL, content TEXT, embedding vector({dim}), metadata JSONB DEFAULT '{{}}', source TEXT, severity TEXT, chain TEXT, created_at TIMESTAMPTZ DEFAULT NOW(), updated_at TIMESTAMPTZ DEFAULT NOW() ); CREATE INDEX IF NOT EXISTS idx_rag_collection ON {VECTOR_TABLE} (collection); CREATE INDEX IF NOT EXISTS idx_rag_metadata ON {VECTOR_TABLE} USING GIN (metadata); CREATE INDEX IF NOT EXISTS idx_rag_content_fts ON {VECTOR_TABLE} USING GIN (to_tsvector('english', COALESCE(content, ''))); """ await self._rpc("exec_sql", {"sql": create_sql}) self._table_ready = True except Exception: pass # Try direct insert to check if table exists if not self._table_ready: try: health_dim = TABLE_DIM test_row = { "id": "_pgvector_health_check", "collection": "_system", "embedding": [0.0] * health_dim, "metadata": json.dumps({"health_check": True}), "source": "system", "severity": "info", } url = f"{_get_url()}/rest/v1/rag_vectors" params = {"on_conflict": "id"} async with httpx.AsyncClient(timeout=10) as client: resp = await client.post(url, json=test_row, params=params, headers=_get_headers()) if resp.status_code in (200, 201): self._table_ready = True logger.info("pgvector: rag_vectors table ready") # Cleanup test row await client.delete( f"{_get_url()}/rest/v1/rag_vectors?id=eq._pgvector_health_check", headers=_get_headers(), ) except Exception as e: logger.warning( "pgvector table not found. Run: /root/backend/supabase_pgvector_setup.sql in Supabase SQL Editor" ) logger.warning(f"Falling back to Redis vector store. Error: {e}") if not self._table_ready: logger.warning("pgvector unavailable - using Redis fallback for vector search") self._table_ready = False # Signal to use Redis fallback return self._table_ready async def insert( self, doc_id: str, collection: str, embedding: list[float], content: str = "", metadata: dict | None = None, source: str = "", severity: str = "medium", chain: str = "", ) -> bool: """Insert a single vector document. Falls back to Redis if pgvector unavailable.""" if not self._table_ready: # Fallback to Redis try: from app.rag_service import _get_redis r = await _get_redis() doc = { "id": doc_id, "collection": collection, "vector": embedding, "dims": len(embedding), "metadata": metadata or {}, "content": content[:5000], "source": source, "severity": severity, } await r.setex(f"rag:{collection}:{doc_id}", 86400 * 30, json.dumps(doc)) await r.sadd(f"rag:idx:{collection}", doc_id) return True except Exception as e: logger.error(f"Redis fallback insert failed: {e}") return False try: row = { "id": doc_id, "collection": collection, "content": content[:10000], "embedding": pad_vector(embedding, TABLE_DIM), "metadata": json.dumps(metadata or {}), "source": source, "severity": severity, "chain": chain, "updated_at": datetime.now(UTC).isoformat(), } await self._upsert(VECTOR_TABLE, [row]) return True except Exception as e: logger.error(f"pgvector insert failed: {e}") return False async def insert_batch(self, docs: list[dict]) -> int: """Insert multiple documents in one batch (max 100 per call).""" count = 0 for i in range(0, len(docs), 100): batch = docs[i : i + 100] rows = [] for doc in batch: rows.append( { "id": doc["id"], "collection": doc["collection"], "content": doc.get("content", "")[:10000], "embedding": pad_vector(doc["embedding"], TABLE_DIM), "metadata": json.dumps(doc.get("metadata", {})), "source": doc.get("source", ""), "severity": doc.get("severity", "medium"), "chain": doc.get("chain", ""), "updated_at": datetime.now(UTC).isoformat(), } ) try: await self._upsert(VECTOR_TABLE, rows) count += len(rows) except Exception as e: logger.error(f"Batch insert failed at batch {i}: {e}") return count async def search( self, query_embedding: list[float], collection: str | None = None, limit: int = 10, min_similarity: float = 0.6, filters: dict | None = None, ) -> list[dict[str, Any]]: """ ANN search using pgvector. Falls back to Redis. """ # Fallback to Redis if not self._table_ready: try: from app.crypto_embeddings import get_embedder from app.rag_service import _get_redis r = await _get_redis() doc_ids = await r.smembers(f"rag:idx:{collection or 'known_scams'}") if not doc_ids: return [] keys = [f"rag:{collection or 'known_scams'}:{did}" for did in doc_ids] pipe = r.pipeline() for k in keys: pipe.get(k) results = await pipe.execute() embedder = await get_embedder() # Use singleton, not CryptoEmbedder() q_len = len(query_embedding) scored = [] for data in results: if not data: continue try: doc = json.loads(data) except Exception: continue dv = doc.get("vector", []) if not dv: continue compare_len = min(q_len, len(dv)) sim = embedder.cosine_similarity(query_embedding[:compare_len], dv[:compare_len]) if sim >= min_similarity: scored.append( { "id": doc["id"], "similarity": round(sim, 4), "content": doc.get("content", ""), "metadata": doc.get("metadata", {}), "source": doc.get("source", ""), "severity": doc.get("severity", ""), } ) scored.sort(key=lambda x: x["similarity"], reverse=True) return scored[:limit] except Exception as e: logger.error(f"Redis search fallback failed: {e}") return [] # pgvector search (primary) # Pad query vector to the table dimension before calling RPC padded_query = pad_vector(query_embedding, TABLE_DIM) try: results = await self._rpc( "search_embeddings", { "query_embedding": padded_query, "namespace": collection or "default", "match_count": limit, "similarity_threshold": min_similarity, }, ) # search_embeddings returns [] for no matches - that's a valid result if results is not None: return [ { "id": r.get("id") or r.get("document_id"), "similarity": round(r.get("similarity", 0), 4), "content": r.get("content", ""), "metadata": r.get("metadata", {}), "source": r.get("source", ""), "severity": r.get("severity", ""), } for r in (results or []) ] except Exception as e: logger.warning(f"search_embeddings RPC failed: {e}") # Fallback: direct SQL query via REST # Build filter query filter_params = {} if collection: filter_params["collection"] = f"eq.{collection}" if filters: for k, v in filters.items(): filter_params[k] = f"eq.{v}" # Construct the query with vector similarity # Using cosine distance: 1 - (embedding <=> query) # Pad or truncate to match the table column dimension embedding_padded = pad_vector(query_embedding, TABLE_DIM) # Format embedding as pgvector literal - safe since it's all floats embedding_literal = "[" + ",".join(f"{x:.8f}" for x in embedding_padded) + "]" # Escape single quotes in collection name to prevent SQL injection safe_collection = collection.replace("'", "''") if collection else "" safe_min_sim = max(0.0, min(1.0, float(min_similarity))) # clamp to valid range try: # Try to use a raw SQL query via RPC collection_filter = f"AND collection = '{safe_collection}'" if safe_collection else "" sql = f""" SELECT id, collection, content, metadata, source, severity, chain, 1 - (embedding <=> '{embedding_literal}'::vector) AS similarity FROM {VECTOR_TABLE} WHERE 1 - (embedding <=> '{embedding_literal}'::vector) > {safe_min_sim} {collection_filter} ORDER BY embedding <=> '{embedding_literal}'::vector LIMIT {int(limit)} """ results = await self._rpc("exec_sql_returning", {"sql": sql}) if results: return [ { "id": r["id"], "similarity": round(r["similarity"], 4), "content": r.get("content", ""), "metadata": r.get("metadata", {}), "source": r.get("source", ""), "severity": r.get("severity", ""), } for r in results ] except Exception as e: logger.warning(f"Direct SQL vector search failed: {e}") return [] async def hybrid_search( self, query_text: str, query_embedding: list[float], collection: str | None = None, limit: int = 10, vector_weight: float = 0.7, ) -> list[dict[str, Any]]: """ Hybrid search: dense vector + sparse text (BM25 via tsvector). Combines semantic similarity with exact keyword matching. Critical for: finding exact code snippets, function names, error messages. """ # Get vector results vector_results = await self.search(query_embedding, collection=collection, limit=limit * 2) # Get text search results text_results = [] try: # Escape user input for PostgreSQL text search - prevent SQL injection safe_query_text = query_text.replace("'", "''").replace("\\", "\\\\") safe_collection = collection.replace("'", "''") if collection else "" collection_filter = f"AND collection = '{safe_collection}'" if safe_collection else "" safe_limit = int(limit * 2) sql = f""" SELECT id, collection, content, metadata, source, severity, ts_rank(to_tsvector('english', COALESCE(content, '')), plainto_tsquery('english', '{safe_query_text}')) AS text_score FROM {VECTOR_TABLE} WHERE to_tsvector('english', COALESCE(content, '')) @@ plainto_tsquery('english', '{safe_query_text}') {collection_filter} ORDER BY text_score DESC LIMIT {safe_limit} """ text_results_raw = await self._rpc("exec_sql_returning", {"sql": sql}) if text_results_raw: max_score = max(r["text_score"] for r in text_results_raw) or 1.0 text_results = [ { "id": r["id"], "similarity": round(r["text_score"] / max_score, 4), "content": r.get("content", ""), "metadata": r.get("metadata", {}), "source": r.get("source", ""), "severity": r.get("severity", ""), } for r in text_results_raw ] except Exception as e: logger.warning(f"Text search failed: {e}") # Merge with weighted Reciprocal Rank Fusion merged = {} for rank, r in enumerate(vector_results): merged[r["id"]] = { **r, "rrf_score": vector_weight / (60 + rank + 1), "match_type": "vector", } for rank, r in enumerate(text_results): rrf = (1 - vector_weight) / (60 + rank + 1) if r["id"] in merged: merged[r["id"]]["rrf_score"] += rrf merged[r["id"]]["match_type"] = "hybrid" else: merged[r["id"]] = { **r, "rrf_score": rrf, "match_type": "text", } # Sort and return top results sorted_results = sorted(merged.values(), key=lambda x: x["rrf_score"], reverse=True) final = [] for r in sorted_results[:limit]: r.pop("rrf_score", None) final.append(r) return final async def get_collection_stats(self) -> dict[str, int]: """Get document counts per collection.""" try: sql = f""" SELECT collection, COUNT(*) as count FROM {VECTOR_TABLE} GROUP BY collection ORDER BY count DESC """ rows = await self._rpc("exec_sql_returning", {"sql": sql}) return {r["collection"]: r["count"] for r in (rows or [])} except Exception: return {} async def delete_collection(self, collection: str) -> int: """Delete all documents in a collection.""" safe_collection = collection.replace("'", "''") try: sql = f"DELETE FROM {VECTOR_TABLE} WHERE collection = '{safe_collection}'" result = await self._rpc("exec_sql", {"sql": sql}) if isinstance(result, list) and result: return result[0].get("count", 1) if isinstance(result[0], dict) else 1 return 1 except Exception: return 0 async def delete_by_id(self, doc_id: str) -> bool: """Delete a single document.""" safe_id = doc_id.replace("'", "''") try: sql = f"DELETE FROM {VECTOR_TABLE} WHERE id = '{safe_id}'" await self._rpc("exec_sql", {"sql": sql}) return True except Exception: return False async def total_docs(self) -> int: """Total document count.""" try: sql = f"SELECT COUNT(*) as count FROM {VECTOR_TABLE}" rows = await self._rpc("exec_sql_returning", {"sql": sql}) return rows[0]["count"] if rows else 0 except Exception: return 0 async def build_hnsw_index( self, m: int = 16, ef_construction: int = 200, ) -> bool: """ Build/rebuild HNSW index for ANN search. HNSW provides better recall than IVFFlat at similar query speeds, and does not require a training step. Parameters: m: Max connections per layer (default 16; higher = more recall, more memory) ef_construction: Build-time search width (default 200; higher = better index quality) """ try: # Drop existing HNSW index if present (to rebuild with new params) try: drop_sql = "DROP INDEX IF EXISTS idx_rag_embedding_hnsw;" await self._rpc("exec_sql", {"sql": drop_sql}) except Exception: pass sql = f""" CREATE INDEX IF NOT EXISTS idx_rag_embedding_hnsw ON {VECTOR_TABLE} USING hnsw (embedding vector_cosine_ops) WITH (m = {int(m)}, ef_construction = {int(ef_construction)}); """ await self._rpc("exec_sql", {"sql": sql}) logger.info(f"HNSW index built (m={m}, ef_construction={ef_construction})") return True except Exception as e: logger.warning(f"HNSW index build failed: {e}") return False async def build_index(self) -> bool: """ Build/rebuild ANN index for fast vector search. Prefers HNSW (higher recall, no training) over IVFFlat. Falls back to IVFFlat if HNSW fails. """ # Try HNSW first (better quality, no training step) hnsw_ok = await self.build_hnsw_index(m=16, ef_construction=200) if hnsw_ok: return True # Fallback: IVFFlat try: sql = f""" CREATE INDEX IF NOT EXISTS idx_rag_embedding_ivfflat ON {VECTOR_TABLE} USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100); """ await self._rpc("exec_sql", {"sql": sql}) logger.info("IVFFlat index built (HNSW fallback)") return True except Exception as e: logger.warning(f"Index build failed (both HNSW and IVFFlat): {e}") return False async def list_distinct(self, column: str, collection: str | None = None) -> list[str]: """Get distinct values for a column.""" # Whitelist allowed columns to prevent SQL injection allowed_columns = {"source", "severity", "chain", "collection"} safe_column = column if column in allowed_columns else "source" try: where = "" if collection: safe_collection = collection.replace("'", "''") where = f"WHERE collection = '{safe_collection}'" sql = f"SELECT DISTINCT {safe_column} FROM {VECTOR_TABLE} {where} ORDER BY {safe_column}" rows = await self._rpc("exec_sql_returning", {"sql": sql}) return [r[safe_column] for r in (rows or [])] except Exception: return [] # ══════════════════════════════════════════════════════════════════════ # SINGLETON # ══════════════════════════════════════════════════════════════════════ _vector_store: SupabaseVectorStore | None = None async def get_vector_store() -> SupabaseVectorStore: global _vector_store # Re-init only if not ready (preserves working singleton across calls) if _vector_store is None or not _vector_store._table_ready: _vector_store = SupabaseVectorStore() await _vector_store.initialize() return _vector_store