langgraph-api 0.0.26__py3-none-any.whl → 0.0.28rc1__py3-none-any.whl

langgraph_storage/queue.py

@@ -1,108 +1,131 @@
 import asyncio
-import logging
-from collections import defaultdict
-from dataclasses import dataclass
-from uuid import UUID
 
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class Message:
-    topic: bytes
-    data: bytes
-
-
-class ContextQueue(asyncio.Queue):
-    """Queue that supports async context manager protocol"""
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: object | None,
-    ) -> None:
-        # Clear the queue
-        while not self.empty():
-            try:
-                self.get_nowait()
-            except asyncio.QueueEmpty:
-                break
-
-
-class StreamManager:
-    def __init__(self):
-        self.queues = defaultdict(list)  # Dict[UUID, List[asyncio.Queue]]
-        self.control_queues = defaultdict(list)
-
-    def get_queues(self, run_id: UUID) -> list[asyncio.Queue]:
-        return self.queues[run_id]
-
-    async def put(self, run_id: UUID, message: Message) -> None:
-        topic = message.topic.decode()
-        if "control" in topic:
-            self.control_queues[run_id].append(message)
-        queues = self.queues.get(run_id, [])
-        coros = [queue.put(message) for queue in queues]
-        results = await asyncio.gather(*coros, return_exceptions=True)
-        for result in results:
-            if isinstance(result, Exception):
-                logger.exception(f"Failed to put message in queue: {result}")
-
-    async def add_queue(self, run_id: UUID) -> asyncio.Queue:
-        queue = ContextQueue()
-        self.queues[run_id].append(queue)
-        for control_msg in self.control_queues[run_id]:
-            try:
-                await queue.put(control_msg)
-            except Exception:
+import structlog
+
+from langgraph_api.config import (
+    BG_JOB_HEARTBEAT,
+    N_JOBS_PER_WORKER,
+    STATS_INTERVAL_SECS,
+)
+from langgraph_api.graph import is_js_graph
+from langgraph_api.schema import Run
+from langgraph_api.webhook import call_webhook
+from langgraph_api.worker import WorkerResult, worker
+from langgraph_storage.database import connect
+from langgraph_storage.ops import Runs
+
+logger = structlog.stdlib.get_logger(__name__)
+
+WORKERS: set[asyncio.Task] = set()
+SHUTDOWN_GRACE_PERIOD_SECS = 5
+
+
+async def queue():
+    concurrency = N_JOBS_PER_WORKER
+    loop = asyncio.get_running_loop()
+    last_stats_secs: int | None = None
+    last_sweep_secs: int | None = None
+    semaphore = asyncio.Semaphore(concurrency)
+    WEBHOOKS: set[asyncio.Task] = set()
+
+    def cleanup(task: asyncio.Task):
+        WORKERS.remove(task)
+        semaphore.release()
+        try:
+            if task.cancelled():
+                return
+            exc = task.exception()
+            if exc and not isinstance(exc, asyncio.CancelledError):
                 logger.exception(
-                    f"Failed to put control message in queue: {control_msg}"
+                    f"Background worker failed for task {task}", exc_info=exc
                 )
-
-        return queue
-
-    async def remove_queue(self, run_id: UUID, queue: asyncio.Queue):
-        if run_id in self.queues:
-            self.queues[run_id].remove(queue)
-            if not self.queues[run_id]:
-                del self.queues[run_id]
-
-
-# Global instance
-stream_manager = StreamManager()
-
-
-async def start_queue() -> None:
-    """Initialize the queue system.
-    In this in-memory implementation, we just need to ensure we have a clean StreamManager instance.
-    """
-    global stream_manager
-    stream_manager = StreamManager()
-
-
-async def stop_queue() -> None:
-    """Clean up the queue system.
-    Clear all queues and stored control messages."""
-    global stream_manager
-
-    # Send 'done' message to all active queues before clearing
-    for run_id in list(stream_manager.queues.keys()):
-        control_message = Message(topic=f"run:{run_id}:control".encode(), data=b"done")
-        for queue in stream_manager.queues[run_id]:
+                return
+            result: WorkerResult | None = task.result()
+            if result and result["webhook"]:
+                hook_task = loop.create_task(
+                    call_webhook(result),
+                    name=f"webhook-{result['run']['run_id']}",
+                )
+                WEBHOOKS.add(hook_task)
+                hook_task.add_done_callback(WEBHOOKS.remove)
+        except asyncio.CancelledError:
+            pass
+        except Exception as exc:
+            logger.exception("Background worker cleanup failed", exc_info=exc)
+
+    await logger.ainfo(f"Starting {concurrency} background workers")
+    try:
+        run: Run | None = None
+        while True:
             try:
-                await queue.put(control_message)
-            except (Exception, RuntimeError):
-                pass  # Ignore errors during shutdown
-
-    # Clear all stored data
-    stream_manager.queues.clear()
-    stream_manager.control_queues.clear()
-
-
-def get_stream_manager() -> StreamManager:
-    """Get the global stream manager instance."""
-    return stream_manager
+                # check if we need to sweep runs
+                do_sweep = (
+                    last_sweep_secs is None
+                    or loop.time() - last_sweep_secs > BG_JOB_HEARTBEAT * 2
+                )
+                # check if we need to update stats
+                if calc_stats := (
+                    last_stats_secs is None
+                    or loop.time() - last_stats_secs > STATS_INTERVAL_SECS
+                ):
+                    last_stats_secs = loop.time()
+                    active = len(WORKERS)
+                    await logger.ainfo(
+                        "Worker stats",
+                        max=concurrency,
+                        available=concurrency - active,
+                        active=active,
+                    )
+                # wait for semaphore to respect concurrency
+                await semaphore.acquire()
+                # skip the wait, if 1st time, or got a run last time
+                wait = run is None and last_stats_secs is not None
+                # try to get a run, handle it
+                run = None
+                async for run, attempt in Runs.next(wait=wait, limit=1):
+                    graph_id = (
+                        run["kwargs"]
+                        .get("config", {})
+                        .get("configurable", {})
+                        .get("graph_id")
+                    )
+
+                    if graph_id and is_js_graph(graph_id):
+                        task_name = f"js-run-{run['run_id']}-attempt-{attempt}"
+                    else:
+                        task_name = f"run-{run['run_id']}-attempt-{attempt}"
+                    task = asyncio.create_task(
+                        worker(run, attempt, loop),
+                        name=task_name,
+                    )
+                    task.add_done_callback(cleanup)
+                    WORKERS.add(task)
+                else:
+                    semaphore.release()
+                # run stats and sweep if needed
+                if calc_stats or do_sweep:
+                    async with connect() as conn:
+                        # update stats if needed
+                        if calc_stats:
+                            stats = await Runs.stats(conn)
+                            await logger.ainfo("Queue stats", **stats)
+                        # sweep runs if needed
+                        if do_sweep:
+                            last_sweep_secs = loop.time()
+                            run_ids = await Runs.sweep(conn)
+                            logger.info("Sweeped runs", run_ids=run_ids)
+            except Exception as exc:
+                # keep trying to run the scheduler indefinitely
+                logger.exception("Background worker scheduler failed", exc_info=exc)
+                semaphore.release()
+                await exit.aclose()
+    finally:
+        logger.info("Shutting down background workers")
+        for task in WORKERS:
+            task.cancel()
+        for task in WEBHOOKS:
+            task.cancel()
+        await asyncio.wait_for(
+            asyncio.gather(*WORKERS, *WEBHOOKS, return_exceptions=True),
+            SHUTDOWN_GRACE_PERIOD_SECS,
+        )

langgraph_storage/retry.py

@@ -7,7 +7,11 @@ P = ParamSpec("P")
 T = TypeVar("T")
 
 
-RETRIABLE_EXCEPTIONS: tuple[type[BaseException], ...] = ()
+class RetryableException(Exception):
+    pass
+
+
+RETRIABLE_EXCEPTIONS: tuple[type[BaseException], ...] = (RetryableException,)
 OVERLOADED_EXCEPTIONS: tuple[type[BaseException], ...] = ()
 
 

langgraph_storage/store.py

@@ -1,4 +1,5 @@
 import os
+import threading
 from collections import defaultdict
 from collections.abc import Iterable
 from typing import Any
@@ -65,6 +66,7 @@ _STORE_FILE = os.path.join(".langgraph_api", "store.pckl")
 _VECTOR_FILE = os.path.join(".langgraph_api", "store.vectors.pckl")
 os.makedirs(".langgraph_api", exist_ok=True)
 STORE = DiskBackedInMemStore()
+BATCHED_STORE = threading.local()
 
 
 def set_store_config(config: dict) -> None:
@@ -77,4 +79,6 @@ def set_store_config(config: dict) -> None:
 
 
 def Store(*args: Any, **kwargs: Any) -> DiskBackedInMemStore:
-    return BatchedStore(STORE)
+    if not hasattr(BATCHED_STORE, "store"):
+        BATCHED_STORE.store = BatchedStore(STORE)
+    return BATCHED_STORE.store

openapi.json

@@ -2962,7 +2962,7 @@
             "type": "string",
             "maxLength": 65536,
             "minLength": 1,
-            "format": "uri",
+            "format": "uri-reference",
             "title": "Webhook",
             "description": "Webhook to call after LangGraph API call is done."
           },
@@ -3380,7 +3380,7 @@
             "type": "string",
             "maxLength": 65536,
             "minLength": 1,
-            "format": "uri",
+            "format": "uri-reference",
             "title": "Webhook",
             "description": "Webhook to call after LangGraph API call is done."
           },
@@ -3607,7 +3607,7 @@
             "type": "string",
             "maxLength": 65536,
             "minLength": 1,
-            "format": "uri",
+            "format": "uri-reference",
             "title": "Webhook",
             "description": "Webhook to call after LangGraph API call is done."
           },