Spaces:
Runtime error
Runtime error
| """ | |
| DocuLens - Supabase persistence layer. | |
| Handles document storage, pipeline run tracking, and stats aggregation. | |
| Falls back gracefully when Supabase is not configured (persist_results=False). | |
| """ | |
| import os | |
| import logging | |
| from datetime import datetime, timedelta, timezone | |
| from typing import Optional | |
| from uuid import uuid4 | |
| logger = logging.getLogger(__name__) | |
| # --------------------------------------------------------------------------- | |
| # Supabase client singleton | |
| # --------------------------------------------------------------------------- | |
| _client = None | |
| SUPABASE_URL = os.environ.get("SUPABASE_URL", "") | |
| SUPABASE_KEY = os.environ.get("SUPABASE_KEY", "") | |
| def get_client(): | |
| """Lazy-init Supabase client. Returns None if not configured.""" | |
| global _client | |
| if _client is not None: | |
| return _client | |
| if not SUPABASE_URL or not SUPABASE_KEY: | |
| logger.info("Supabase not configured — persistence disabled") | |
| return None | |
| try: | |
| from supabase import create_client | |
| _client = create_client(SUPABASE_URL, SUPABASE_KEY) | |
| logger.info("Supabase client initialized") | |
| return _client | |
| except Exception as e: | |
| logger.warning(f"Failed to init Supabase: {e}") | |
| return None | |
| def is_enabled() -> bool: | |
| """Check if persistence is available.""" | |
| return get_client() is not None | |
| # --------------------------------------------------------------------------- | |
| # Documents | |
| # --------------------------------------------------------------------------- | |
| def save_document( | |
| filename: str, | |
| file_type: str, | |
| file_size: int, | |
| page_count: int = 1, | |
| org_id: Optional[str] = None, | |
| user_id: Optional[str] = None, | |
| ) -> Optional[str]: | |
| """Save document metadata. Returns document_id or None.""" | |
| client = get_client() | |
| if not client: | |
| return None | |
| doc_id = str(uuid4()) | |
| try: | |
| client.table("documents").insert({ | |
| "id": doc_id, | |
| "org_id": org_id, | |
| "user_id": user_id, | |
| "filename": filename, | |
| "file_type": file_type, | |
| "file_size": file_size, | |
| "page_count": page_count, | |
| }).execute() | |
| return doc_id | |
| except Exception as e: | |
| logger.warning(f"Failed to save document: {e}") | |
| return None | |
| # --------------------------------------------------------------------------- | |
| # Pipeline Runs | |
| # --------------------------------------------------------------------------- | |
| def create_run( | |
| document_id: Optional[str], | |
| use_case: str = "invoice", | |
| org_id: Optional[str] = None, | |
| user_id: Optional[str] = None, | |
| ) -> Optional[str]: | |
| """Create a pipeline run. Returns run_id or None.""" | |
| client = get_client() | |
| if not client: | |
| return None | |
| run_id = str(uuid4()) | |
| try: | |
| client.table("pipeline_runs").insert({ | |
| "id": run_id, | |
| "org_id": org_id, | |
| "user_id": user_id, | |
| "document_id": document_id, | |
| "use_case": use_case, | |
| "status": "running", | |
| "started_at": datetime.now(timezone.utc).isoformat(), | |
| }).execute() | |
| return run_id | |
| except Exception as e: | |
| logger.warning(f"Failed to create run: {e}") | |
| return None | |
| def complete_run( | |
| run_id: str, | |
| status: str, | |
| overall_result: str, | |
| extraction_data: dict, | |
| processing_time_ms: int, | |
| error_message: Optional[str] = None, | |
| ) -> bool: | |
| """Mark a run as completed/failed with its results.""" | |
| client = get_client() | |
| if not client or not run_id: | |
| return False | |
| try: | |
| client.table("pipeline_runs").update({ | |
| "status": status, | |
| "overall_result": overall_result, | |
| "extraction_data": extraction_data, | |
| "processing_time_ms": processing_time_ms, | |
| "error_message": error_message, | |
| "completed_at": datetime.now(timezone.utc).isoformat(), | |
| }).eq("id", run_id).execute() | |
| return True | |
| except Exception as e: | |
| logger.warning(f"Failed to complete run: {e}") | |
| return False | |
| def save_stage_result( | |
| run_id: str, | |
| stage_type: str, | |
| stage_name: str, | |
| status: str, | |
| sort_order: int, | |
| output: Optional[dict] = None, | |
| error_message: Optional[str] = None, | |
| duration_ms: Optional[int] = None, | |
| ) -> bool: | |
| """Save a single stage result within a run.""" | |
| client = get_client() | |
| if not client or not run_id: | |
| return False | |
| try: | |
| client.table("stage_results").insert({ | |
| "id": str(uuid4()), | |
| "run_id": run_id, | |
| "stage_type": stage_type, | |
| "stage_name": stage_name, | |
| "status": status, | |
| "sort_order": sort_order, | |
| "output": output, | |
| "error_message": error_message, | |
| "duration_ms": duration_ms, | |
| "started_at": datetime.now(timezone.utc).isoformat(), | |
| "completed_at": datetime.now(timezone.utc).isoformat(), | |
| }).execute() | |
| return True | |
| except Exception as e: | |
| logger.warning(f"Failed to save stage result: {e}") | |
| return False | |
| # --------------------------------------------------------------------------- | |
| # Queries — Runs listing + stats | |
| # --------------------------------------------------------------------------- | |
| def list_runs( | |
| org_id: Optional[str] = None, | |
| user_id: Optional[str] = None, | |
| status: Optional[str] = None, | |
| result: Optional[str] = None, | |
| use_case: Optional[str] = None, | |
| search: Optional[str] = None, | |
| page: int = 1, | |
| per_page: int = 25, | |
| ) -> dict: | |
| """List pipeline runs with filtering and pagination.""" | |
| client = get_client() | |
| if not client: | |
| return {"runs": [], "pagination": {"page": 1, "per_page": per_page, "total": 0, "pages": 0}} | |
| try: | |
| # Select only the columns needed for the list view — deliberately | |
| # exclude extraction_data which contains large base64 page images. | |
| list_columns = ( | |
| "id, document_id, use_case, status, overall_result, " | |
| "processing_time_ms, started_at, completed_at, error_message" | |
| ) | |
| # Use !inner join when searching by filename so PostgREST can | |
| # filter on the related table; regular join otherwise. | |
| if search: | |
| doc_join = "documents!inner(filename, file_type, file_size)" | |
| else: | |
| doc_join = "documents(filename, file_type, file_size)" | |
| query = client.table("pipeline_runs").select( | |
| f"{list_columns}, {doc_join}", | |
| count="exact", | |
| ) | |
| if user_id: | |
| query = query.eq("user_id", user_id) | |
| if org_id: | |
| query = query.eq("org_id", org_id) | |
| if status: | |
| query = query.eq("status", status) | |
| if result: | |
| query = query.eq("overall_result", result) | |
| if use_case: | |
| query = query.eq("use_case", use_case) | |
| if search: | |
| query = query.ilike("documents.filename", f"%{search}%") | |
| # Pagination | |
| offset = (page - 1) * per_page | |
| query = query.order("started_at", desc=True).range(offset, offset + per_page - 1) | |
| response = query.execute() | |
| total = response.count or 0 | |
| # Batch-fetch stage_results for all runs in one query instead of N+1 | |
| run_ids = [row["id"] for row in response.data] | |
| stages_by_run: dict[str, list] = {rid: [] for rid in run_ids} | |
| if run_ids: | |
| stages_resp = client.table("stage_results").select( | |
| "run_id, stage_type, stage_name, status, duration_ms, sort_order" | |
| ).in_("run_id", run_ids).order("sort_order").execute() | |
| for s in stages_resp.data: | |
| stages_by_run.setdefault(s["run_id"], []).append(s) | |
| runs = [] | |
| for row in response.data: | |
| doc = row.get("documents") or {} | |
| runs.append({ | |
| "id": row["id"], | |
| "document": { | |
| "id": row.get("document_id"), | |
| "filename": doc.get("filename", "unknown"), | |
| "file_type": doc.get("file_type", ""), | |
| }, | |
| "use_case": row.get("use_case", "invoice"), | |
| "status": row.get("status", "pending"), | |
| "overall_result": row.get("overall_result"), | |
| "stages": [ | |
| { | |
| "type": s["stage_type"], | |
| "name": s["stage_name"], | |
| "status": s["status"], | |
| "duration_ms": s.get("duration_ms"), | |
| } | |
| for s in stages_by_run.get(row["id"], []) | |
| ], | |
| "processing_time_ms": row.get("processing_time_ms"), | |
| "started_at": row.get("started_at"), | |
| "completed_at": row.get("completed_at"), | |
| "error_message": row.get("error_message"), | |
| }) | |
| return { | |
| "runs": runs, | |
| "pagination": { | |
| "page": page, | |
| "per_page": per_page, | |
| "total": total, | |
| "pages": max(1, (total + per_page - 1) // per_page), | |
| }, | |
| } | |
| except Exception as e: | |
| logger.warning(f"Failed to list runs: {e}") | |
| return {"runs": [], "pagination": {"page": 1, "per_page": per_page, "total": 0, "pages": 0}} | |
| def get_run(run_id: str) -> Optional[dict]: | |
| """Get a single run with full details.""" | |
| client = get_client() | |
| if not client: | |
| return None | |
| try: | |
| response = client.table("pipeline_runs").select( | |
| "*, documents(filename, file_type, file_size)" | |
| ).eq("id", run_id).single().execute() | |
| row = response.data | |
| if not row: | |
| return None | |
| stages_resp = client.table("stage_results").select("*").eq( | |
| "run_id", run_id | |
| ).order("sort_order").execute() | |
| doc = row.get("documents") or {} | |
| return { | |
| "id": row["id"], | |
| "document": { | |
| "id": row.get("document_id"), | |
| "filename": doc.get("filename", "unknown"), | |
| "file_type": doc.get("file_type", ""), | |
| }, | |
| "use_case": row.get("use_case", "invoice"), | |
| "status": row.get("status"), | |
| "overall_result": row.get("overall_result"), | |
| "extraction_data": row.get("extraction_data"), | |
| "stages": [ | |
| { | |
| "type": s["stage_type"], | |
| "name": s["stage_name"], | |
| "status": s["status"], | |
| "output": s.get("output"), | |
| "error_message": s.get("error_message"), | |
| "duration_ms": s.get("duration_ms"), | |
| } | |
| for s in stages_resp.data | |
| ], | |
| "processing_time_ms": row.get("processing_time_ms"), | |
| "started_at": row.get("started_at"), | |
| "completed_at": row.get("completed_at"), | |
| } | |
| except Exception as e: | |
| logger.warning(f"Failed to get run: {e}") | |
| return None | |
| def get_stats( | |
| org_id: Optional[str] = None, | |
| user_id: Optional[str] = None, | |
| period_days: int = 30, | |
| ) -> dict: | |
| """Get aggregate stats for dashboard.""" | |
| client = get_client() | |
| if not client: | |
| return _empty_stats(period_days) | |
| try: | |
| since = (datetime.now(timezone.utc) - timedelta(days=period_days)).isoformat() | |
| query = client.table("pipeline_runs").select("*").gte("started_at", since) | |
| if user_id: | |
| query = query.eq("user_id", user_id) | |
| if org_id: | |
| query = query.eq("org_id", org_id) | |
| response = query.execute() | |
| runs = response.data or [] | |
| # Aggregate totals | |
| total = len(runs) | |
| passed = sum(1 for r in runs if r.get("overall_result") == "passed") | |
| failed = sum(1 for r in runs if r.get("overall_result") == "failed") | |
| needs_review = sum(1 for r in runs if r.get("overall_result") == "needs_review") | |
| times = [r.get("processing_time_ms", 0) for r in runs if r.get("processing_time_ms")] | |
| avg_ms = int(sum(times) / len(times)) if times else 0 | |
| # Group by day | |
| by_day: dict[str, dict] = {} | |
| for r in runs: | |
| day = r.get("started_at", "")[:10] | |
| if day not in by_day: | |
| by_day[day] = {"date": day, "passed": 0, "failed": 0, "needs_review": 0} | |
| result = r.get("overall_result", "") | |
| if result in by_day[day]: | |
| by_day[day][result] += 1 | |
| # Group by use case | |
| by_uc: dict[str, int] = {} | |
| for r in runs: | |
| uc = r.get("use_case", "invoice") | |
| by_uc[uc] = by_uc.get(uc, 0) + 1 | |
| end_date = datetime.now(timezone.utc).strftime("%Y-%m-%d") | |
| start_date = (datetime.now(timezone.utc) - timedelta(days=period_days)).strftime("%Y-%m-%d") | |
| return { | |
| "period": {"start": start_date, "end": end_date}, | |
| "totals": { | |
| "runs": total, | |
| "passed": passed, | |
| "failed": failed, | |
| "needs_review": needs_review, | |
| "avg_processing_ms": avg_ms, | |
| }, | |
| "by_day": sorted(by_day.values(), key=lambda d: d["date"]), | |
| "by_use_case": [ | |
| {"slug": slug, "name": slug.replace("-", " ").title(), "count": count} | |
| for slug, count in sorted(by_uc.items(), key=lambda x: -x[1]) | |
| ], | |
| } | |
| except Exception as e: | |
| logger.warning(f"Failed to get stats: {e}") | |
| return _empty_stats(period_days) | |
| def _empty_stats(period_days: int) -> dict: | |
| end_date = datetime.now(timezone.utc).strftime("%Y-%m-%d") | |
| start_date = (datetime.now(timezone.utc) - timedelta(days=period_days)).strftime("%Y-%m-%d") | |
| return { | |
| "period": {"start": start_date, "end": end_date}, | |
| "totals": {"runs": 0, "passed": 0, "failed": 0, "needs_review": 0, "avg_processing_ms": 0}, | |
| "by_day": [], | |
| "by_use_case": [], | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Golden Copies — user-verified extraction ground truth | |
| # --------------------------------------------------------------------------- | |
| def save_golden_copy( | |
| user_id: str, | |
| filename: str, | |
| extracted_fields: dict, | |
| line_items: list, | |
| use_case: str = "invoice", | |
| document_id: Optional[str] = None, | |
| run_id: Optional[str] = None, | |
| org_id: Optional[str] = None, | |
| notes: str = "", | |
| ) -> Optional[str]: | |
| """Save a user-verified extraction as a golden copy. Returns the golden copy id.""" | |
| client = get_client() | |
| if not client: | |
| return None | |
| try: | |
| row = { | |
| "user_id": user_id, | |
| "filename": filename, | |
| "extracted_fields": extracted_fields, | |
| "line_items": line_items, | |
| "use_case": use_case, | |
| "notes": notes, | |
| } | |
| if document_id: | |
| row["document_id"] = document_id | |
| if run_id: | |
| row["run_id"] = run_id | |
| if org_id: | |
| row["org_id"] = org_id | |
| resp = client.table("golden_copies").insert(row).execute() | |
| return resp.data[0]["id"] if resp.data else None | |
| except Exception as e: | |
| logger.warning(f"Failed to save golden copy: {e}") | |
| return None | |
| def list_golden_copies( | |
| user_id: str, | |
| use_case: Optional[str] = None, | |
| page: int = 1, | |
| per_page: int = 20, | |
| ) -> dict: | |
| """List golden copies for a user.""" | |
| client = get_client() | |
| if not client: | |
| return {"golden_copies": [], "total": 0} | |
| try: | |
| query = ( | |
| client.table("golden_copies") | |
| .select("id, filename, use_case, notes, created_at, updated_at", count="exact") | |
| .eq("user_id", user_id) | |
| .order("created_at", desc=True) | |
| .range((page - 1) * per_page, page * per_page - 1) | |
| ) | |
| if use_case: | |
| query = query.eq("use_case", use_case) | |
| resp = query.execute() | |
| return { | |
| "golden_copies": resp.data or [], | |
| "total": resp.count or 0, | |
| } | |
| except Exception as e: | |
| logger.warning(f"Failed to list golden copies: {e}") | |
| return {"golden_copies": [], "total": 0} | |
| def get_golden_copy(golden_id: str) -> Optional[dict]: | |
| """Retrieve a single golden copy by id.""" | |
| client = get_client() | |
| if not client: | |
| return None | |
| try: | |
| resp = client.table("golden_copies").select("*").eq("id", golden_id).single().execute() | |
| return resp.data | |
| except Exception as e: | |
| logger.warning(f"Failed to get golden copy: {e}") | |
| return None | |
| # --------------------------------------------------------------------------- | |
| # Prompt Versions — CRUD for configurable prompts | |
| # --------------------------------------------------------------------------- | |
| def list_prompt_versions( | |
| use_case: str, | |
| stage: str, | |
| active_only: bool = False, | |
| ) -> list[dict]: | |
| """List prompt versions for a (use_case, stage) pair, newest first.""" | |
| client = get_client() | |
| if not client: | |
| return [] | |
| try: | |
| query = client.table("prompt_versions").select("*") \ | |
| .eq("use_case", use_case) \ | |
| .eq("stage", stage) | |
| if active_only: | |
| query = query.eq("is_active", True) | |
| query = query.order("version", desc=True) | |
| response = query.execute() | |
| return response.data or [] | |
| except Exception as e: | |
| logger.warning(f"Failed to list prompt versions: {e}") | |
| return [] | |
| def get_active_prompt(use_case: str, stage: str) -> Optional[str]: | |
| """Get the currently active prompt text for a (use_case, stage) pair. | |
| Returns None if no active prompt is stored (caller should fall back | |
| to the hardcoded PROMPT_LIBRARY). | |
| """ | |
| client = get_client() | |
| if not client: | |
| return None | |
| try: | |
| response = client.table("prompt_versions").select("prompt_text") \ | |
| .eq("use_case", use_case) \ | |
| .eq("stage", stage) \ | |
| .eq("is_active", True) \ | |
| .order("version", desc=True) \ | |
| .limit(1) \ | |
| .execute() | |
| rows = response.data or [] | |
| if rows: | |
| return rows[0]["prompt_text"] | |
| return None | |
| except Exception as e: | |
| logger.warning(f"Failed to get active prompt: {e}") | |
| return None | |
| def save_prompt_version( | |
| use_case: str, | |
| stage: str, | |
| prompt_text: str, | |
| created_by: Optional[str] = None, | |
| ) -> Optional[dict]: | |
| """Create a new prompt version, making it the active one. | |
| Deactivates any previously active version for the same (use_case, stage). | |
| Returns the new row or None on failure. | |
| """ | |
| client = get_client() | |
| if not client: | |
| return None | |
| try: | |
| # Find current max version number | |
| existing = client.table("prompt_versions").select("version") \ | |
| .eq("use_case", use_case) \ | |
| .eq("stage", stage) \ | |
| .order("version", desc=True) \ | |
| .limit(1) \ | |
| .execute() | |
| next_version = 1 | |
| if existing.data: | |
| next_version = existing.data[0]["version"] + 1 | |
| # Deactivate previous active version(s) | |
| client.table("prompt_versions").update({"is_active": False}) \ | |
| .eq("use_case", use_case) \ | |
| .eq("stage", stage) \ | |
| .eq("is_active", True) \ | |
| .execute() | |
| # Insert new version | |
| row = { | |
| "id": str(uuid4()), | |
| "use_case": use_case, | |
| "stage": stage, | |
| "version": next_version, | |
| "prompt_text": prompt_text, | |
| "is_active": True, | |
| "created_by": created_by, | |
| "created_at": datetime.now(timezone.utc).isoformat(), | |
| } | |
| client.table("prompt_versions").insert(row).execute() | |
| return row | |
| except Exception as e: | |
| logger.warning(f"Failed to save prompt version: {e}") | |
| return None | |
| def activate_prompt_version( | |
| use_case: str, | |
| stage: str, | |
| version: int, | |
| ) -> bool: | |
| """Set a specific version as the active prompt.""" | |
| client = get_client() | |
| if not client: | |
| return False | |
| try: | |
| # Deactivate all for this (use_case, stage) | |
| client.table("prompt_versions").update({"is_active": False}) \ | |
| .eq("use_case", use_case) \ | |
| .eq("stage", stage) \ | |
| .execute() | |
| # Activate the requested version | |
| client.table("prompt_versions").update({"is_active": True}) \ | |
| .eq("use_case", use_case) \ | |
| .eq("stage", stage) \ | |
| .eq("version", version) \ | |
| .execute() | |
| return True | |
| except Exception as e: | |
| logger.warning(f"Failed to activate prompt version: {e}") | |
| return False | |
| def get_all_active_prompts() -> list[dict]: | |
| """Get all active prompts across all use cases and stages.""" | |
| client = get_client() | |
| if not client: | |
| return [] | |
| try: | |
| response = client.table("prompt_versions").select("*") \ | |
| .eq("is_active", True) \ | |
| .order("use_case") \ | |
| .order("stage") \ | |
| .execute() | |
| return response.data or [] | |
| except Exception as e: | |
| logger.warning(f"Failed to get all active prompts: {e}") | |
| return [] | |