RugMunchMedia/rmi-backend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

merge: chore/cleanup-remove-bloat-and-secrets into main

Browse source
This commit is contained in:
Crypto Rug Munch 2026-07-02 01:24:22 +07:00
commit bde2f3a97d
1173 changed files with 437609 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

694
app/supabase_vector.py Normal file
View file

@ -0,0 +1,694 @@
#!/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