fix: use MemorySaver to avoid async checkpoint NotImplementedError

src/ai/graph.py

@@ -14,7 +14,7 @@ from __future__ import annotations
 
 import redis
 import structlog
-from langgraph.checkpoint.redis import RedisSaver
+from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import END, StateGraph
 
 from ..config import settings
@@ -117,9 +117,8 @@ def get_compiled_graph():
     """
     workflow = build_extraction_graph()
 
-    # Redis checkpointer — stores full graph state per thread_id (batch_id)
-    redis_conn = redis.Redis.from_url(settings.REDIS_URL)
-    checkpointer = RedisSaver(redis_client=redis_conn)
+    # Using MemorySaver to support async ainvoke correctly
+    checkpointer = MemorySaver()
 
     graph = workflow.compile(checkpointer=checkpointer, interrupt_before=["human_review"])
 

src/tasks/ocr_tasks.py

@@ -121,7 +121,7 @@ def preprocess_document(self, batch_id: str) -> None:
     log.info("Starting extraction pipeline", batch_id=batch_id)
     try:
         config = {"configurable": {"thread_id": batch_id}}
-        context = asyncio.get_event_loop().run_until_complete(_load_batch_context(batch_id))
+        context = asyncio.run(_load_batch_context(batch_id))
         initial_state = {
             "batch_id": batch_id,
             "company_id": context["batch"].get("company_id") or "",
@@ -135,10 +135,10 @@ def preprocess_document(self, batch_id: str) -> None:
             "steps": [],
         }
         # Run sync wrapper around async graph
-        result = asyncio.get_event_loop().run_until_complete(
+        result = asyncio.run(
             extraction_graph.ainvoke(initial_state, config=config)
         )
-        asyncio.get_event_loop().run_until_complete(_persist_graph_result(batch_id, result))
+        asyncio.run(_persist_graph_result(batch_id, result))
         log.info("Extraction pipeline complete", batch_id=batch_id)
     except Exception as exc:
         log.error("Pipeline failed", batch_id=batch_id, error=str(exc))