Spaces:
Running
Running
| import json | |
| import unittest | |
| import httpx | |
| from backend.storage.supabase_store import SupabaseJobStore | |
| # The 7 tables of the baseline schema. | |
| SCHEMA_TABLES = [ | |
| "profiles", | |
| "submissions", | |
| "reviews", | |
| "feedback", | |
| "pipeline_steps", | |
| "provider_calls", | |
| "evidence_items", | |
| ] | |
| class SupabaseTransport: | |
| def __init__(self): | |
| self.storage = {} | |
| self.tables = {name: {} for name in SCHEMA_TABLES} | |
| self._ids = {name: 0 for name in self.tables} | |
| def __call__(self, request: httpx.Request) -> httpx.Response: | |
| path = request.url.path | |
| method = request.method | |
| if path.startswith("/storage/v1/object/") and method == "POST": | |
| object_path = path.split("/object/", 1)[1] | |
| self.storage[object_path] = request.content | |
| return httpx.Response(200, json={"Key": object_path}) | |
| if path.startswith("/storage/v1/object/") and method == "GET": | |
| object_path = path.split("/object/", 1)[1] | |
| content = self.storage.get(object_path) | |
| if content is None: | |
| return httpx.Response(404, json={"error": "not found"}) | |
| return httpx.Response(200, content=content) | |
| if path.startswith("/rest/v1/rpc/") and method == "POST": | |
| return self._handle_rpc(path.rsplit("/", 1)[-1], request) | |
| if not path.startswith("/rest/v1/"): | |
| return httpx.Response(404, json={"error": f"Unhandled route: {method} {path}"}) | |
| table = path.rsplit("/", 1)[-1] | |
| if table not in self.tables: | |
| return httpx.Response(404, json={"error": f"Unknown table: {table}"}) | |
| if method == "POST": | |
| payload = json.loads(request.content.decode("utf-8")) | |
| on_conflict = request.url.params.get("on_conflict") | |
| inserted = [] | |
| for row in payload: | |
| if table == "profiles" and on_conflict == "email": | |
| existing = next((r for r in self.tables[table].values() if r["email"] == row["email"]), None) | |
| if existing: | |
| existing.update(row) | |
| inserted.append(existing) | |
| continue | |
| if table == "reviews" and on_conflict == "submission_id": | |
| existing = next((r for r in self.tables[table].values() if r["submission_id"] == row["submission_id"]), None) | |
| if existing: | |
| existing.update(row) | |
| inserted.append(existing) | |
| continue | |
| self._ids[table] += 1 | |
| row.setdefault("id", self._ids[table]) | |
| row.setdefault("created_at", "2026-01-01T00:00:00+00:00") | |
| row.setdefault("completed_at", None) | |
| row.setdefault("finished_at", None) | |
| self.tables[table][row["id"]] = row | |
| inserted.append(row) | |
| return httpx.Response(201, json=inserted) | |
| if method == "GET": | |
| rows = list(self.tables[table].values()) | |
| for key, value in request.url.params.multi_items(): | |
| if key in {"select", "limit", "order"} or "." in key: | |
| continue | |
| if value.startswith("eq."): | |
| expected = value[3:] | |
| rows = [r for r in rows if str(r.get(key)) == expected] | |
| if table == "submissions": | |
| for row in rows: | |
| review = next( | |
| (r for r in self.tables["reviews"].values() if r.get("submission_id") == row.get("id")), | |
| None, | |
| ) | |
| row["reviews"] = [review] if review else [] | |
| return httpx.Response(200, json=rows) | |
| if method == "PATCH": | |
| rows = list(self.tables[table].values()) | |
| for key, value in request.url.params.multi_items(): | |
| if value.startswith("eq."): | |
| expected = value[3:] | |
| rows = [r for r in rows if str(r.get(key)) == expected] | |
| payload = json.loads(request.content.decode("utf-8")) | |
| for row in rows: | |
| row.update(payload) | |
| return httpx.Response(200, json=rows) | |
| return httpx.Response(405, json={"error": f"Unhandled method: {method}"}) | |
| def _handle_rpc(self, fn: str, request: httpx.Request) -> httpx.Response: | |
| """Simulate the set_review_completed Postgres RPC (/rest/v1/rpc/<fn>).""" | |
| params = json.loads(request.content.decode("utf-8")) | |
| if fn == "set_review_completed": | |
| submission = next( | |
| (r for r in self.tables["submissions"].values() | |
| if r.get("access_key") == params["p_access_key"]), | |
| None, | |
| ) | |
| if submission is None: | |
| return httpx.Response(400, json={"code": "P0002", "message": "submission not found"}) | |
| review_row = { | |
| "submission_id": submission["id"], | |
| "review_json": params["p_review"], | |
| "overall_score": params["p_overall_score"], | |
| "summary": params["p_summary"], | |
| } | |
| existing = next( | |
| (r for r in self.tables["reviews"].values() | |
| if r.get("submission_id") == submission["id"]), | |
| None, | |
| ) | |
| if existing: | |
| existing.update(review_row) | |
| else: | |
| self._ids["reviews"] += 1 | |
| review_row["id"] = self._ids["reviews"] | |
| review_row.setdefault("created_at", "2026-01-01T00:00:00+00:00") | |
| self.tables["reviews"][review_row["id"]] = review_row | |
| submission["status"] = "completed" | |
| submission["completed_at"] = "2026-01-01T00:00:00+00:00" | |
| submission["error"] = None | |
| return httpx.Response(204) | |
| return httpx.Response(404, json={"error": f"Unknown rpc: {fn}"}) | |
| class SupabaseJobStoreTests(unittest.IsolatedAsyncioTestCase): | |
| def setUp(self): | |
| self.transport = SupabaseTransport() | |
| self.store = SupabaseJobStore( | |
| supabase_url="https://example.supabase.co", | |
| supabase_key="service-role-key", | |
| storage_bucket="paper-review-pdfs", | |
| transport=httpx.MockTransport(self.transport), | |
| ) | |
| async def asyncTearDown(self): | |
| await self.store.aclose() | |
| async def test_create_job_uses_submission_id_as_storage_root(self): | |
| job = await self.store.create_job( | |
| access_key="sk-pm-test", | |
| email="user@example.com", | |
| filename="paper.pdf", | |
| pdf_bytes=b"%PDF-1.4 test", | |
| ) | |
| self.assertEqual(job["status"], "pending") | |
| self.assertEqual(job["access_key"], "sk-pm-test") | |
| self.assertEqual(job["paper_filename"], "paper.pdf") | |
| self.assertNotIn("submission_key", job) | |
| submission = next(iter(self.transport.tables["submissions"].values())) | |
| self.assertTrue(submission["pdf_storage_path"]) | |
| self.assertNotIn("submission_key", submission) | |
| # Paper file metadata is folded into the metadata jsonb (keeps the table narrow). | |
| self.assertEqual(submission["metadata"]["paper_size_bytes"], len(b"%PDF-1.4 test")) | |
| storage_key = next(iter(self.transport.storage.keys())) | |
| self.assertIn(f"submissions/{submission['id']}/input/", storage_key) | |
| async def test_get_pdf_bytes_returns_original_pdf(self): | |
| await self.store.create_job("sk-pm-pdf", "p@example.com", "paper.pdf", b"%PDF-bytes") | |
| result = await self.store.get_pdf_bytes("sk-pm-pdf") | |
| self.assertIsNotNone(result) | |
| data, filename = result | |
| self.assertEqual(data, b"%PDF-bytes") | |
| self.assertEqual(filename, "paper.pdf") | |
| async def test_saves_parsed_markdown_text_review_and_feedback(self): | |
| await self.store.create_job("sk-pm-review", "person@example.com", "draft.pdf", b"%PDF") | |
| parsed = await self.store.save_parsed_markdown( | |
| "sk-pm-review", | |
| markdown="# Paper\nBody", | |
| provider="docling", | |
| parser_version="kaggle-docling", | |
| raw_json={"pages": 2}, | |
| ) | |
| self.assertEqual(parsed["provider"], "docling") | |
| self.assertEqual(parsed["char_count"], 12) | |
| # Markdown text is queryable directly from the DB. | |
| submission = next(iter(self.transport.tables["submissions"].values())) | |
| self.assertEqual(submission["parsed_markdown_text"], "# Paper\nBody") | |
| self.assertEqual(await self.store.get_parsed_markdown("sk-pm-review"), "# Paper\nBody") | |
| review = { | |
| "paper_summary": "Summary", | |
| "overall_assessment": 3.5, | |
| "soundness": 4, | |
| "confidence": 3, | |
| } | |
| await self.store.set_completed("sk-pm-review", review) | |
| fetched = await self.store.get_job("sk-pm-review") | |
| self.assertEqual(fetched["status"], "completed") | |
| self.assertEqual(fetched["review"]["overall_assessment"], 3.5) | |
| self.assertTrue(fetched["has_parsed_markdown"]) | |
| # review_json is the canonical store; the denormalized `scores` column was dropped. | |
| saved_review = next(iter(self.transport.tables["reviews"].values())) | |
| self.assertEqual(saved_review["review_json"]["overall_assessment"], 3.5) | |
| self.assertEqual(saved_review["review_json"]["soundness"], 4) | |
| self.assertNotIn("scores", saved_review) | |
| feedback = await self.store.save_feedback( | |
| "sk-pm-review", | |
| { | |
| "helpfulness": "high", | |
| "has_critical_error": False, | |
| "has_suggestions": True, | |
| "comments": "Useful", | |
| "submitted_at": "2026-06-16T00:00:00+00:00", | |
| }, | |
| ) | |
| self.assertEqual(feedback["comments"], "Useful") | |
| self.assertEqual(feedback["metadata"]["submitted_at"], "2026-06-16T00:00:00+00:00") | |
| async def test_records_step_output_costs_and_evidence(self): | |
| await self.store.create_job("sk-pm-usage", "cost@example.com", "cost.pdf", b"%PDF") | |
| step = await self.store.start_pipeline_step( | |
| "sk-pm-usage", | |
| "generate_search_queries", | |
| step_order=4, | |
| ) | |
| await self.store.record_llm_call( | |
| access_key="sk-pm-usage", | |
| step_id=step["id"], | |
| step_name="generate_search_queries", | |
| provider="openai", | |
| model="gpt-4o-mini", | |
| response_format="json", | |
| prompt_tokens=1000, | |
| completion_tokens=500, | |
| total_tokens=1500, | |
| input_price_per_1m=0.15, | |
| output_price_per_1m=0.60, | |
| latency_ms=1234, | |
| status="success", | |
| error=None, | |
| ) | |
| await self.store.record_api_call( | |
| access_key="sk-pm-usage", | |
| step_id=step["id"], | |
| step_name="search_related_papers", | |
| provider="tavily", | |
| operation="search", | |
| status_code=200, | |
| latency_ms=321, | |
| request_count=3, | |
| status="success", | |
| error=None, | |
| cost_usd=0.024, | |
| metadata={"credits": 3, "price_per_credit": 0.008}, | |
| ) | |
| await self.store.record_search_queries("sk-pm-usage", step["id"], ["review llm eval"]) | |
| await self.store.finish_pipeline_step( | |
| "sk-pm-usage", step["id"], "completed", | |
| latency_ms=2000, | |
| output_json={"queries": ["review llm eval"]}, | |
| output_summary="1 queries", | |
| ) | |
| # provider_calls: an LLM row (priced by tokens) and a Tavily API row (priced by credits). | |
| calls = list(self.transport.tables["provider_calls"].values()) | |
| self.assertEqual([c["call_type"] for c in calls], ["llm", "api"]) | |
| llm_call = next(c for c in calls if c["call_type"] == "llm") | |
| self.assertAlmostEqual(llm_call["cost_usd"], 0.00045, places=8) | |
| self.assertEqual(llm_call["metadata"]["prompt_tokens"], 1000) | |
| api_call = next(c for c in calls if c["call_type"] == "api") | |
| self.assertAlmostEqual(api_call["cost_usd"], 0.024, places=8) | |
| self.assertEqual(api_call["metadata"]["credits"], 3) | |
| # step output persisted on the pipeline_steps row | |
| saved_step = self.transport.tables["pipeline_steps"][step["id"]] | |
| self.assertEqual(saved_step["status"], "completed") | |
| self.assertEqual(saved_step["output_json"], {"queries": ["review llm eval"]}) | |
| # evidence linked to the step | |
| evidence = next(iter(self.transport.tables["evidence_items"].values())) | |
| self.assertEqual(evidence["evidence_type"], "search_query") | |
| self.assertEqual(evidence["query_text"], "review llm eval") | |
| self.assertEqual(evidence["pipeline_step_id"], step["id"]) | |