import asyncio import re from dataclasses import dataclass import json import os import time from typing import Any, Dict, List, Optional from sql_security import strip_trailing_semicolon try: import asyncpg except Exception: # pragma: no cover - optional until requirements are installed asyncpg = None def _get_database_url() -> Optional[str]: return ( os.environ.get("READONLY_DATABASE_URL") or os.environ.get("SUPABASE_DB_URL") or os.environ.get("DATABASE_URL") or os.environ.get("POSTGRES_URL") ) @dataclass class QueryExecution: success: bool rows: List[Dict[str, Any]] columns: List[str] row_count: int execution_ms: float execution_plan: Any = None indexes_used: Optional[List[str]] = None error: Optional[str] = None class PostgresService: def __init__(self) -> None: self.database_url = _get_database_url() self.pool = None @property def configured(self) -> bool: return bool(self.database_url and asyncpg is not None) async def connect(self) -> None: if not self.configured or self.pool is not None: return self.pool = await asyncpg.create_pool( dsn=self.database_url, min_size=1, max_size=int(os.environ.get("POSTGRES_POOL_SIZE", "5")), command_timeout=12, ) async def close(self) -> None: if self.pool is not None: await self.pool.close() self.pool = None async def execute_select(self, sql: str, limit: int = 100, timeout: float = 10.0) -> QueryExecution: if self.pool is None: raise RuntimeError("PostgreSQL is not configured") safe_sql = strip_trailing_semicolon(sql) limited_sql = f"SELECT * FROM ({safe_sql}) AS neuralvault_query LIMIT {int(limit)}" explain_sql = f"EXPLAIN (FORMAT JSON, ANALYZE false) {safe_sql}" start = time.perf_counter() try: async with self.pool.acquire() as conn: # ELEGANT Spec 5: Sandboxing statement timeouts at transaction session level async with conn.transaction(): await conn.execute("SET statement_timeout = 2000; SET lock_timeout = 1000;") rows = await conn.fetch(limited_sql, timeout=timeout) plan = await conn.fetchval(explain_sql, timeout=timeout) except Exception as exc: return QueryExecution( success=False, rows=[], columns=[], row_count=0, execution_ms=round((time.perf_counter() - start) * 1000, 2), error=str(exc), ) dict_rows = [dict(row) for row in rows] columns = list(dict_rows[0].keys()) if dict_rows else [] return QueryExecution( success=True, rows=dict_rows, columns=columns, row_count=len(dict_rows), execution_ms=round((time.perf_counter() - start) * 1000, 2), execution_plan=plan, indexes_used=_extract_indexes(plan), ) async def health(self) -> dict: if asyncpg is None: return {"configured": False, "ok": False, "reason": "asyncpg is not installed"} if not self.database_url: return {"configured": False, "ok": False, "reason": "No PostgreSQL URL configured"} if self.pool is None: return {"configured": True, "ok": False, "reason": "Pool not initialized"} try: async with self.pool.acquire() as conn: await conn.fetchval("SELECT 1") vector_count = await _safe_fetchval(conn, "SELECT COUNT(*) FROM products WHERE embedding IS NOT NULL") hnsw_indexes = await conn.fetch( """ SELECT indexname FROM pg_indexes WHERE schemaname = 'public' AND indexdef ILIKE '%hnsw%' """ ) return { "configured": True, "ok": True, "vector_count": vector_count, "hnsw_indexes": [row["indexname"] for row in hnsw_indexes], } except Exception as exc: return {"configured": True, "ok": False, "reason": str(exc)} async def run_startup_validation(self) -> dict: if self.pool is None: return {"ok": False, "error": "Database not connected"} try: async with self.pool.acquire() as conn: # 1. Check pgvector extension has_vector = await conn.fetchval( "SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'vector')" ) # 2. Check products table has_products = await conn.fetchval( "SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_schema = 'public' AND table_name = 'products')" ) # 3. Check HNSW index hnsw_rows = await conn.fetch( "SELECT indexname FROM pg_indexes WHERE schemaname = 'public' AND tablename = 'products' AND indexdef ILIKE '%hnsw%'" ) has_hnsw = len(hnsw_rows) > 0 return { "ok": True, "pgvector_extension": has_vector, "products_table": has_products, "hnsw_index": has_hnsw } except Exception as exc: return {"ok": False, "error": str(exc)} async def hybrid_search( self, query: str, embedding: List[float], mode: str = "hybrid", limit: int = 10, ef_search: int = 40, price_limit: Optional[float] = None, ) -> dict: if self.pool is None: raise RuntimeError("PostgreSQL is not configured") vector_literal = "[" + ",".join(f"{value:.8f}" for value in embedding) + "]" timings = {"embedding_ms": 0.0, "vector_ms": 0.0, "fts_ms": 0.0, "fusion_ms": 0.0} vector_rows: List[Dict[str, Any]] = [] text_rows: List[Dict[str, Any]] = [] async with self.pool.acquire() as conn: await conn.execute(f"SET hnsw.ef_search = {int(ef_search)}") if mode in ("vector", "hybrid"): start = time.perf_counter() if price_limit is not None: # Spec 7: Vector Metadata Pre-Filtering using HNSW indexes with WHERE constraints vector_rows = [ dict(row) for row in await conn.fetch( """ SELECT id::text, title AS name, category, price::text, rating, 1 - (embedding <=> $1::vector) AS vector_score FROM products WHERE embedding IS NOT NULL AND price <= $2 ORDER BY embedding <=> $1::vector LIMIT 50 """, vector_literal, float(price_limit), ) ] else: vector_rows = [ dict(row) for row in await conn.fetch( """ SELECT id::text, title AS name, category, price::text, rating, 1 - (embedding <=> $1::vector) AS vector_score FROM products WHERE embedding IS NOT NULL ORDER BY embedding <=> $1::vector LIMIT 50 """, vector_literal, ) ] timings["vector_ms"] = round((time.perf_counter() - start) * 1000, 2) if mode in ("fulltext", "hybrid"): start = time.perf_counter() if price_limit is not None: text_rows = [ dict(row) for row in await conn.fetch( """ SELECT id::text, title AS name, category, price::text, rating, ts_rank( to_tsvector('english', title || ' ' || COALESCE(description, '')), plainto_tsquery('english', $1) ) AS text_score FROM products WHERE to_tsvector('english', title || ' ' || COALESCE(description, '')) @@ plainto_tsquery('english', $1) AND price <= $2 ORDER BY text_score DESC LIMIT 50 """, query, float(price_limit), ) ] else: text_rows = [ dict(row) for row in await conn.fetch( """ SELECT id::text, title AS name, category, price::text, rating, ts_rank( to_tsvector('english', title || ' ' || COALESCE(description, '')), plainto_tsquery('english', $1) ) AS text_score FROM products WHERE to_tsvector('english', title || ' ' || COALESCE(description, '')) @@ plainto_tsquery('english', $1) ORDER BY text_score DESC LIMIT 50 """, query, ) ] timings["fts_ms"] = round((time.perf_counter() - start) * 1000, 2) start = time.perf_counter() # Stage 1: Blended candidate results via RRF Fusion candidates = _merge_search_results(vector_rows, text_rows, mode, limit=50) # Spec 8 Stage 2: Cross-Encoder Reranking if mode == "hybrid" and candidates: try: from embedding_service import rerank_candidates results = await rerank_candidates(query, candidates, limit=limit) except Exception as exc: print(f"Failed to run Cross-Encoder reranker: {exc}") results = candidates[:limit] else: results = candidates[:limit] timings["fusion_ms"] = round((time.perf_counter() - start) * 1000, 2) timings["total_ms"] = round(sum(timings.values()), 2) return { "results": results, "timings": timings, "ef_search": ef_search, "mode": mode, "price_limit": price_limit, } async def _safe_fetchval(conn, sql: str) -> Any: try: return await conn.fetchval(sql) except Exception: return None def _merge_search_results(vector_rows: List[dict], text_rows: List[dict], mode: str, limit: int) -> List[dict]: if mode == "vector": return [_format_search_row(row, vector_rank=i + 1) for i, row in enumerate(vector_rows[:limit])] if mode == "fulltext": return [_format_search_row(row, text_rank=i + 1) for i, row in enumerate(text_rows[:limit])] by_id: Dict[str, dict] = {} for rank, row in enumerate(vector_rows, start=1): item = by_id.setdefault(row["id"], row.copy()) item["vector_rank"] = rank item["vector_score"] = float(row.get("vector_score") or 0) for rank, row in enumerate(text_rows, start=1): item = by_id.setdefault(row["id"], row.copy()) item["text_rank"] = rank item["text_score"] = float(row.get("text_score") or 0) fused = [] for item in by_id.values(): vector_rank = item.get("vector_rank") text_rank = item.get("text_rank") rrf_score = 0.0 if vector_rank: rrf_score += 1 / (60 + vector_rank) if text_rank: rrf_score += 1 / (60 + text_rank) formatted = _format_search_row(item, vector_rank=vector_rank, text_rank=text_rank) formatted["rrf_score"] = round(rrf_score, 4) fused.append(formatted) return sorted(fused, key=lambda row: row["rrf_score"], reverse=True)[:limit] def _format_search_row(row: dict, vector_rank: Optional[int] = None, text_rank: Optional[int] = None) -> dict: vector_score = float(row.get("vector_score") or 0) text_score = float(row.get("text_score") or 0) rrf_score = row.get("rrf_score") if rrf_score is None: ranks = [rank for rank in (vector_rank, text_rank) if rank] rrf_score = sum(1 / (60 + rank) for rank in ranks) if ranks else 0 reasons = [] if vector_rank: reasons.append(f"semantic rank #{vector_rank}") if text_rank: reasons.append(f"full-text rank #{text_rank}") return { "id": row.get("id"), "name": row.get("name"), "category": row.get("category"), "price": row.get("price") or "n/a", "rating": row.get("rating"), "vector_score": round(vector_score, 4), "text_score": round(text_score, 4), "rrf_score": round(float(rrf_score), 4), "why_matched": "Matched by " + " and ".join(reasons) if reasons else "Matched product search query", "in_stock": True, } def _extract_indexes(plan: Any) -> List[str]: if not plan: return [] if isinstance(plan, str): try: plan = json.loads(plan) except Exception: return sorted(set(re.findall(r"Index Scan using ([a-zA-Z0-9_]+)", plan))) indexes = set() def walk(node: Any) -> None: if isinstance(node, list): for item in node: walk(item) elif isinstance(node, dict): if "Index Name" in node: indexes.add(node["Index Name"]) for value in node.values(): walk(value) walk(plan) return sorted(indexes) postgres = PostgresService()