Retry AI qualifier failures once

src/vk_parser_app/workers/ai_qualifier.py

@@ -277,20 +277,40 @@ class AIQualifierWorker:
             usage.get("estimated_cost_usd"),
         )
 
-    async def mark_posts_failed(self, post_ids: list[int], error: str) -> None:
+    async def mark_posts_after_failed_batch(self, post_ids: list[int], error: str, max_attempts: int = 2) -> tuple[list[int], list[int]]:
         if not post_ids:
-            return
-        await self.pool.execute(
+            return [], []
+        rows = await self.pool.fetch(
             """
+            WITH target AS (
+              SELECT unnest($1::bigint[]) AS id
+            ),
+            attempts AS (
+              SELECT t.id, COUNT(b.id) AS failed_batches
+              FROM target t
+              LEFT JOIN ai_qualification_batches b
+                ON b.status='failed'
+               AND b.request_json @> jsonb_build_array(jsonb_build_object('id', t.id))
+              GROUP BY t.id
+            )
             UPDATE raw_posts
-            SET qualification_status='failed',
-                qualification_reason=$2,
+            SET qualification_status=CASE WHEN attempts.failed_batches < $3 THEN 'pending' ELSE 'failed' END,
+                qualification_reason=CASE
+                  WHEN attempts.failed_batches < $3 THEN $2 || ' (retry once)'
+                  ELSE $2
+                END,
                 updated_at=NOW()
-            WHERE id=ANY($1::bigint[])
+            FROM attempts
+            WHERE raw_posts.id=attempts.id
+            RETURNING raw_posts.id, raw_posts.qualification_status
             """,
             post_ids,
-            error[:1000],
+            error[:900],
+            max_attempts,
         )
+        retry_ids = [int(row["id"]) for row in rows if row["qualification_status"] == "pending"]
+        failed_ids = [int(row["id"]) for row in rows if row["qualification_status"] == "failed"]
+        return retry_ids, failed_ids
 
     async def report_failed_batch(self, model: str, post_ids: list[int], error: str) -> None:
         meta = {"error": error[:300], "post_ids": post_ids, "model": model}
@@ -419,9 +439,17 @@ class AIQualifierWorker:
             logger.info("AI qualifier batch done: id={} posts={} usage={}", batch_id, len(results), usage)
             return True
         except Exception as exc:
-            await self.mark_posts_failed(post_ids, str(exc))
             await self.finish_batch(batch_id, "failed", error=str(exc))
-            await self.report_failed_batch(normalize_model(provider, model), post_ids, str(exc))
+            retry_ids, failed_ids = await self.mark_posts_after_failed_batch(post_ids, str(exc))
+            if failed_ids:
+                await self.report_failed_batch(normalize_model(provider, model), failed_ids, str(exc))
+            else:
+                await self.heartbeat.beat(
+                    self.pool,
+                    status="retrying_after_error",
+                    meta={"error": str(exc)[:300], "post_ids": retry_ids, "model": normalize_model(provider, model)},
+                    force=True,
+                )
             logger.exception("AI qualifier batch failed: id={} error={}", batch_id, exc)
             return True