Spaces:
Sleeping
Sleeping
Commit ·
8b949ea
1
Parent(s): 604f6ac
fix(ocr): fix event loop initialization in celery and update gemini model versions
Browse files- src/config.py +2 -2
- src/tasks/celery_app.py +13 -3
- src/tasks/ocr_tasks.py +2 -1
src/config.py
CHANGED
|
@@ -102,8 +102,8 @@ class Settings(BaseSettings):
|
|
| 102 |
|
| 103 |
# ── AI / LLM ─────────────────────────────────────────────────────────────
|
| 104 |
GEMINI_API_KEY: SecretStr = Field(..., description="Google Gemini API key")
|
| 105 |
-
GEMINI_MODEL_PRIMARY: str = "gemini-
|
| 106 |
-
GEMINI_MODEL_FALLBACK: str = "gemini-
|
| 107 |
OPENAI_API_KEY: SecretStr = "" # type: ignore[assignment]
|
| 108 |
EMBEDDING_MODEL: str = "text-embedding-3-small"
|
| 109 |
|
|
|
|
| 102 |
|
| 103 |
# ── AI / LLM ─────────────────────────────────────────────────────────────
|
| 104 |
GEMINI_API_KEY: SecretStr = Field(..., description="Google Gemini API key")
|
| 105 |
+
GEMINI_MODEL_PRIMARY: str = "gemini-1.5-pro"
|
| 106 |
+
GEMINI_MODEL_FALLBACK: str = "gemini-1.5-flash"
|
| 107 |
OPENAI_API_KEY: SecretStr = "" # type: ignore[assignment]
|
| 108 |
EMBEDDING_MODEL: str = "text-embedding-3-small"
|
| 109 |
|
src/tasks/celery_app.py
CHANGED
|
@@ -42,12 +42,22 @@ 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 |
-
|
| 49 |
-
asyncio.set_event_loop(
|
| 50 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 51 |
|
| 52 |
# ── Celery app ────────────────────────────────────────────────────────────────
|
| 53 |
celery_app = Celery(
|
|
|
|
| 42 |
from celery.signals import worker_process_init
|
| 43 |
import asyncio
|
| 44 |
|
| 45 |
+
_worker_loop = None
|
| 46 |
+
|
| 47 |
@worker_process_init.connect
|
| 48 |
def init_celery_worker(**kwargs):
|
| 49 |
+
global _worker_loop
|
| 50 |
from ..dependencies import init_supabase
|
| 51 |
+
_worker_loop = asyncio.new_event_loop()
|
| 52 |
+
asyncio.set_event_loop(_worker_loop)
|
| 53 |
+
_worker_loop.run_until_complete(init_supabase())
|
| 54 |
+
|
| 55 |
+
def get_worker_loop():
|
| 56 |
+
global _worker_loop
|
| 57 |
+
if _worker_loop is None:
|
| 58 |
+
_worker_loop = asyncio.new_event_loop()
|
| 59 |
+
asyncio.set_event_loop(_worker_loop)
|
| 60 |
+
return _worker_loop
|
| 61 |
|
| 62 |
# ── Celery app ────────────────────────────────────────────────────────────────
|
| 63 |
celery_app = Celery(
|
src/tasks/ocr_tasks.py
CHANGED
|
@@ -120,7 +120,8 @@ 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 |
-
|
|
|
|
| 124 |
config = {"configurable": {"thread_id": batch_id}}
|
| 125 |
context = loop.run_until_complete(_load_batch_context(batch_id))
|
| 126 |
initial_state = {
|
|
|
|
| 120 |
from ..ai.graph import extraction_graph
|
| 121 |
log.info("Starting extraction pipeline", batch_id=batch_id)
|
| 122 |
try:
|
| 123 |
+
from .celery_app import get_worker_loop
|
| 124 |
+
loop = get_worker_loop()
|
| 125 |
config = {"configurable": {"thread_id": batch_id}}
|
| 126 |
context = loop.run_until_complete(_load_batch_context(batch_id))
|
| 127 |
initial_state = {
|