| """Supabase-backed review store.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import Any |
| from uuid import uuid4 |
|
|
| from supabase import Client |
|
|
| from app.enums import CatSlug, EvaluatorLabel |
| from app.exceptions import DuplicateFeedbackError, DuplicateSessionError |
| from app.models import CandidateImage, EvaluationTask, Invite |
|
|
|
|
| def _is_unique_violation(exc: Exception) -> bool: |
| message = str(exc).lower() |
| return "duplicate key" in message or "23505" in message |
|
|
|
|
| class SupabaseReviewStore: |
| def __init__(self, client: Client) -> None: |
| self._client = client |
|
|
| def get_invite_by_token(self, token: str) -> Invite | None: |
| response = ( |
| self._client.table("invites") |
| .select("id, token, evaluator_label, is_active") |
| .eq("token", token) |
| .limit(1) |
| .execute() |
| ) |
| rows = response.data or [] |
| if not rows: |
| return None |
| row = rows[0] |
| if not row.get("is_active"): |
| return None |
| return Invite( |
| id=row["id"], |
| token=row["token"], |
| evaluator_label=EvaluatorLabel(row["evaluator_label"]), |
| is_active=row["is_active"], |
| ) |
|
|
| def get_or_create_session(self, invite_id: str) -> str: |
| response = ( |
| self._client.table("sessions") |
| .select("id") |
| .eq("invite_id", invite_id) |
| .limit(1) |
| .execute() |
| ) |
| rows = response.data or [] |
| if rows: |
| return rows[0]["id"] |
|
|
| session_id = str(uuid4()) |
| try: |
| self._client.table("sessions").insert( |
| {"id": session_id, "invite_id": invite_id} |
| ).execute() |
| except Exception as exc: |
| if not _is_unique_violation(exc): |
| raise |
| response = ( |
| self._client.table("sessions") |
| .select("id") |
| .eq("invite_id", invite_id) |
| .limit(1) |
| .execute() |
| ) |
| rows = response.data or [] |
| if not rows: |
| raise DuplicateSessionError from exc |
| return rows[0]["id"] |
| return session_id |
|
|
| def _completed_task_ids(self, session_id: str) -> set[str]: |
| response = ( |
| self._client.table("feedback") |
| .select("task_id") |
| .eq("session_id", session_id) |
| .execute() |
| ) |
| return {row["task_id"] for row in (response.data or []) if row.get("task_id")} |
|
|
| def get_completed_task_ids(self, session_id: str) -> set[str]: |
| return self._completed_task_ids(session_id) |
|
|
| def _build_evaluation_task(self, task_row: dict[str, Any]) -> EvaluationTask: |
| cat = task_row.get("cats") or {} |
| candidate_response = ( |
| self._client.table("task_candidates") |
| .select( |
| "slot, images(id, local_path, storage_url, prompt, base_model, lora_version, seed)" |
| ) |
| .eq("task_id", task_row["id"]) |
| .order("slot") |
| .execute() |
| ) |
| candidate_rows = candidate_response.data or [] |
|
|
| candidates: list[CandidateImage] = [] |
| for row in candidate_rows: |
| image = row.get("images") or {} |
| image_id = image.get("id") |
| if image_id is None: |
| continue |
| candidates.append( |
| CandidateImage( |
| id=image_id, |
| slot=row["slot"], |
| local_path=image.get("local_path") or "", |
| storage_url=image.get("storage_url"), |
| prompt=image.get("prompt"), |
| base_model=image.get("base_model"), |
| lora_version=image.get("lora_version"), |
| seed=image.get("seed"), |
| ) |
| ) |
|
|
| cat_slug_raw = cat.get("slug") or "" |
| return EvaluationTask( |
| id=task_row["id"], |
| test_id=task_row["test_id"], |
| batch_id=task_row["batch_id"], |
| cat_slug=CatSlug(cat_slug_raw), |
| cat_display_name=cat.get("display_name", ""), |
| candidates=candidates, |
| ) |
|
|
| def _list_task_rows(self) -> list[dict[str, Any]]: |
| response = ( |
| self._client.table("evaluation_tasks") |
| .select("id, test_id, batch_id, cat_id, cats(slug, display_name)") |
| .order("created_at") |
| .execute() |
| ) |
| return response.data or [] |
|
|
| def count_evaluation_tasks(self) -> int: |
| response = ( |
| self._client.table("evaluation_tasks") |
| .select("id", count="exact") |
| .execute() |
| ) |
| return response.count or 0 |
|
|
| def count_completed_tasks(self, session_id: str) -> int: |
| return len(self._completed_task_ids(session_id)) |
|
|
| def get_next_evaluation_task(self, session_id: str) -> EvaluationTask | None: |
| completed = self._completed_task_ids(session_id) |
| for task_row in self._list_task_rows(): |
| if task_row["id"] not in completed: |
| return self._build_evaluation_task(task_row) |
| return None |
|
|
| def get_evaluation_task(self, task_id: str) -> EvaluationTask | None: |
| response = ( |
| self._client.table("evaluation_tasks") |
| .select("id, test_id, batch_id, cat_id, cats(slug, display_name)") |
| .eq("id", task_id) |
| .limit(1) |
| .execute() |
| ) |
| rows = response.data or [] |
| if not rows: |
| return None |
| return self._build_evaluation_task(rows[0]) |
|
|
| def list_all_evaluation_tasks(self) -> list[EvaluationTask]: |
| return [self._build_evaluation_task(row) for row in self._list_task_rows()] |
|
|
| def list_completed_tasks(self, session_id: str) -> list[EvaluationTask]: |
| completed = self._completed_task_ids(session_id) |
| tasks: list[EvaluationTask] = [] |
| for task_row in self._list_task_rows(): |
| if task_row["id"] in completed: |
| tasks.append(self._build_evaluation_task(task_row)) |
| return tasks |
|
|
| def get_task_round_number(self, task_id: str) -> int: |
| for index, task_row in enumerate(self._list_task_rows(), start=1): |
| if task_row["id"] == task_id: |
| return index |
| return 1 |
|
|
| def get_accepted_image_ids(self, session_id: str, task_id: str) -> set[str]: |
| response = ( |
| self._client.table("feedback") |
| .select("image_id, accepted") |
| .eq("session_id", session_id) |
| .eq("task_id", task_id) |
| .execute() |
| ) |
| return { |
| row["image_id"] |
| for row in (response.data or []) |
| if row.get("accepted") == 1 and row.get("image_id") |
| } |
|
|
| def get_first_evaluation_task(self) -> EvaluationTask | None: |
| rows = self._list_task_rows() |
| if not rows: |
| return None |
| return self._build_evaluation_task(rows[0]) |
|
|
| def has_feedback_for_session_task(self, session_id: str, task_id: str) -> bool: |
| response = ( |
| self._client.table("feedback") |
| .select("id") |
| .eq("session_id", session_id) |
| .eq("task_id", task_id) |
| .limit(1) |
| .execute() |
| ) |
| return bool(response.data) |
|
|
| def submit_feedback( |
| self, |
| session_id: str, |
| task_id: str, |
| selections: dict[str, bool], |
| ) -> None: |
| rows: list[dict[str, Any]] = [] |
| for image_id, accepted in selections.items(): |
| rows.append( |
| { |
| "session_id": session_id, |
| "task_id": task_id, |
| "image_id": image_id, |
| "accepted": 1 if accepted else 0, |
| } |
| ) |
| try: |
| self._client.table("feedback").upsert( |
| rows, |
| on_conflict="session_id,task_id,image_id", |
| ).execute() |
| except Exception as exc: |
| if _is_unique_violation(exc): |
| raise DuplicateFeedbackError from exc |
| raise |
|
|
|
|
| def get_review_store() -> SupabaseReviewStore: |
| from app.config import create_supabase_client |
|
|
| return SupabaseReviewStore(create_supabase_client()) |
|
|
