Spaces:
Sleeping
Sleeping
Commit ·
518aad9
1
Parent(s): 7a16f4e
fix(worker): initialize supabase client on celery worker startup
Browse files- src/tasks/celery_app.py +9 -0
- src/tasks/ocr_tasks.py +4 -3
src/tasks/celery_app.py
CHANGED
|
@@ -39,6 +39,15 @@ except Exception: # pragma: no cover - provide lightweight fallbacks for tests
|
|
| 39 |
pass
|
| 40 |
|
| 41 |
from ..config import settings
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 42 |
|
| 43 |
# ── Celery app ────────────────────────────────────────────────────────────────
|
| 44 |
celery_app = Celery(
|
|
|
|
| 39 |
pass
|
| 40 |
|
| 41 |
from ..config import settings
|
| 42 |
+
from celery.signals import worker_process_init
|
| 43 |
+
import asyncio
|
| 44 |
+
|
| 45 |
+
@worker_process_init.connect
|
| 46 |
+
def init_celery_worker(**kwargs):
|
| 47 |
+
from ..dependencies import init_supabase
|
| 48 |
+
loop = asyncio.new_event_loop()
|
| 49 |
+
asyncio.set_event_loop(loop)
|
| 50 |
+
loop.run_until_complete(init_supabase())
|
| 51 |
|
| 52 |
# ── Celery app ────────────────────────────────────────────────────────────────
|
| 53 |
celery_app = Celery(
|
src/tasks/ocr_tasks.py
CHANGED
|
@@ -120,8 +120,9 @@ def preprocess_document(self, batch_id: str) -> None:
|
|
| 120 |
from ..ai.graph import extraction_graph
|
| 121 |
log.info("Starting extraction pipeline", batch_id=batch_id)
|
| 122 |
try:
|
|
|
|
| 123 |
config = {"configurable": {"thread_id": batch_id}}
|
| 124 |
-
context =
|
| 125 |
initial_state = {
|
| 126 |
"batch_id": batch_id,
|
| 127 |
"company_id": context["batch"].get("company_id") or "",
|
|
@@ -135,10 +136,10 @@ def preprocess_document(self, batch_id: str) -> None:
|
|
| 135 |
"steps": [],
|
| 136 |
}
|
| 137 |
# Run sync wrapper around async graph
|
| 138 |
-
result =
|
| 139 |
extraction_graph.ainvoke(initial_state, config=config)
|
| 140 |
)
|
| 141 |
-
|
| 142 |
log.info("Extraction pipeline complete", batch_id=batch_id)
|
| 143 |
except Exception as exc:
|
| 144 |
log.error("Pipeline failed", batch_id=batch_id, error=str(exc))
|
|
|
|
| 120 |
from ..ai.graph import extraction_graph
|
| 121 |
log.info("Starting extraction pipeline", batch_id=batch_id)
|
| 122 |
try:
|
| 123 |
+
loop = asyncio.get_event_loop()
|
| 124 |
config = {"configurable": {"thread_id": batch_id}}
|
| 125 |
+
context = loop.run_until_complete(_load_batch_context(batch_id))
|
| 126 |
initial_state = {
|
| 127 |
"batch_id": batch_id,
|
| 128 |
"company_id": context["batch"].get("company_id") or "",
|
|
|
|
| 136 |
"steps": [],
|
| 137 |
}
|
| 138 |
# Run sync wrapper around async graph
|
| 139 |
+
result = loop.run_until_complete(
|
| 140 |
extraction_graph.ainvoke(initial_state, config=config)
|
| 141 |
)
|
| 142 |
+
loop.run_until_complete(_persist_graph_result(batch_id, result))
|
| 143 |
log.info("Extraction pipeline complete", batch_id=batch_id)
|
| 144 |
except Exception as exc:
|
| 145 |
log.error("Pipeline failed", batch_id=batch_id, error=str(exc))
|