langgraph-api 0.0.1__py3-none-any.whl

langgraph_api/queue.py

@@ -0,0 +1,245 @@
+import asyncio
+from contextlib import AsyncExitStack
+from datetime import UTC, datetime
+from random import random
+from typing import cast
+
+import structlog
+from langgraph.pregel.debug import CheckpointPayload
+
+from langgraph_api.config import BG_JOB_NO_DELAY, STATS_INTERVAL_SECS
+from langgraph_api.errors import (
+    UserInterrupt,
+    UserRollback,
+)
+from langgraph_api.http import get_http_client
+from langgraph_api.js.remote import RemoteException
+from langgraph_api.metadata import incr_runs
+from langgraph_api.schema import Run
+from langgraph_api.stream import (
+    astream_state,
+    consume,
+)
+from langgraph_api.utils import AsyncConnectionProto
+from langgraph_storage.database import connect
+from langgraph_storage.ops import Runs, Threads
+from langgraph_storage.retry import RETRIABLE_EXCEPTIONS
+
+logger = structlog.stdlib.get_logger(__name__)
+
+WORKERS: set[asyncio.Task] = set()
+MAX_RETRY_ATTEMPTS = 3
+SHUTDOWN_GRACE_PERIOD_SECS = 5
+
+
+def ms(after: datetime, before: datetime) -> int:
+    return int((after - before).total_seconds() * 1000)
+
+
+async def queue(concurrency: int, timeout: float):
+    loop = asyncio.get_running_loop()
+    last_stats_secs: int | None = None
+    semaphore = asyncio.Semaphore(concurrency)
+
+    def cleanup(task: asyncio.Task):
+        WORKERS.remove(task)
+        semaphore.release()
+        try:
+            if exc := task.exception():
+                logger.exception("Background worker failed", exc_info=exc)
+        except asyncio.CancelledError:
+            pass
+
+    await logger.ainfo(f"Starting {concurrency} background workers")
+    try:
+        while True:
+            try:
+                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,
+                    )
+                await semaphore.acquire()
+                exit = AsyncExitStack()
+                conn = await exit.enter_async_context(connect())
+                if calc_stats:
+                    stats = await Runs.stats(conn)
+                    await logger.ainfo("Queue stats", **stats)
+                if tup := await exit.enter_async_context(Runs.next(conn)):
+                    task = asyncio.create_task(
+                        worker(timeout, exit, conn, *tup),
+                        name=f"run-{tup[0]['run_id']}-attempt-{tup[1]}",
+                    )
+                    task.add_done_callback(cleanup)
+                    WORKERS.add(task)
+                else:
+                    semaphore.release()
+                    await exit.aclose()
+                    await asyncio.sleep(0 if BG_JOB_NO_DELAY else random())
+            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()
+                await asyncio.sleep(0 if BG_JOB_NO_DELAY else random())
+    finally:
+        logger.info("Shutting down background workers")
+        for task in WORKERS:
+            task.cancel()
+        await asyncio.wait_for(
+            asyncio.gather(*WORKERS, return_exceptions=True), SHUTDOWN_GRACE_PERIOD_SECS
+        )
+
+
+async def worker(
+    timeout: float,
+    exit: AsyncExitStack,
+    conn: AsyncConnectionProto,
+    run: Run,
+    attempt: int,
+):
+    run_id = run["run_id"]
+    if attempt == 1:
+        incr_runs()
+    async with Runs.enter(run_id) as done, exit:
+        temporary = run["kwargs"].get("temporary", False)
+        webhook = run["kwargs"].pop("webhook", None)
+        checkpoint: CheckpointPayload | None = None
+        exception: Exception | None = None
+        status: str | None = None
+        run_started_at = datetime.now(UTC)
+        run_created_at = run["created_at"].isoformat()
+        await logger.ainfo(
+            "Starting background run",
+            run_id=str(run_id),
+            run_attempt=attempt,
+            run_created_at=run_created_at,
+            run_started_at=run_started_at.isoformat(),
+            run_queue_ms=ms(run_started_at, run["created_at"]),
+        )
+
+        def on_checkpoint(checkpoint_arg: CheckpointPayload):
+            nonlocal checkpoint
+            checkpoint = checkpoint_arg
+
+        try:
+            if attempt > MAX_RETRY_ATTEMPTS:
+                raise RuntimeError(f"Run {run['run_id']} exceeded max attempts")
+            if temporary:
+                stream = astream_state(
+                    AsyncExitStack(), conn, cast(Run, run), attempt, done
+                )
+            else:
+                stream = astream_state(
+                    AsyncExitStack(),
+                    conn,
+                    cast(Run, run),
+                    attempt,
+                    done,
+                    on_checkpoint=on_checkpoint,
+                )
+            await asyncio.wait_for(consume(stream, run_id), timeout)
+            await logger.ainfo(
+                "Background run succeeded",
+                run_id=str(run_id),
+                run_attempt=attempt,
+                run_created_at=run_created_at,
+                run_started_at=run_started_at.isoformat(),
+                run_ended_at=datetime.now().isoformat(),
+                run_exec_ms=ms(datetime.now(UTC), run_started_at),
+            )
+            status = "success"
+            await Runs.set_status(conn, run_id, "success")
+        except TimeoutError as e:
+            exception = e
+            status = "timeout"
+            await logger.awarning(
+                "Background run timed out",
+                run_id=str(run_id),
+                run_attempt=attempt,
+                run_created_at=run_created_at,
+                run_started_at=run_started_at.isoformat(),
+                run_ended_at=datetime.now().isoformat(),
+                run_exec_ms=ms(datetime.now(UTC), run_started_at),
+            )
+            await Runs.set_status(conn, run_id, "timeout")
+        except UserRollback as e:
+            exception = e
+            status = "rollback"
+            await logger.ainfo(
+                "Background run rolled back",
+                run_id=str(run_id),
+                run_attempt=attempt,
+                run_created_at=run_created_at,
+                run_started_at=run_started_at.isoformat(),
+                run_ended_at=datetime.now().isoformat(),
+                run_exec_ms=ms(datetime.now(UTC), run_started_at),
+            )
+            await Runs.delete(conn, run_id, thread_id=run["thread_id"])
+        except UserInterrupt as e:
+            exception = e
+            status = "interrupted"
+            await logger.ainfo(
+                "Background run interrupted",
+                run_id=str(run_id),
+                run_attempt=attempt,
+                run_created_at=run_created_at,
+                run_started_at=run_started_at.isoformat(),
+                run_ended_at=datetime.now().isoformat(),
+                run_exec_ms=ms(datetime.now(UTC), run_started_at),
+            )
+            await Runs.set_status(conn, run_id, "interrupted")
+        except RETRIABLE_EXCEPTIONS as e:
+            exception = e
+            status = "retry"
+            await logger.awarning(
+                "Background run failed, will retry",
+                exc_info=True,
+                run_id=str(run_id),
+                run_attempt=attempt,
+                run_created_at=run_created_at,
+                run_started_at=run_started_at.isoformat(),
+                run_ended_at=datetime.now().isoformat(),
+                run_exec_ms=ms(datetime.now(UTC), run_started_at),
+            )
+            raise
+            # Note we re-raise here, thus marking the run
+            # as available to be picked up by another worker
+        except Exception as exc:
+            exception = exc
+            status = "error"
+            await logger.aexception(
+                "Background run failed",
+                exc_info=not isinstance(exc, RemoteException),
+                run_id=str(run_id),
+                run_attempt=attempt,
+                run_created_at=run_created_at,
+                run_started_at=run_started_at.isoformat(),
+                run_ended_at=datetime.now().isoformat(),
+                run_exec_ms=ms(datetime.now(UTC), run_started_at),
+            )
+            await Runs.set_status(conn, run_id, "error")
+        # delete or set status of thread
+        if temporary:
+            await Threads.delete(conn, run["thread_id"])
+        else:
+            await Threads.set_status(conn, run["thread_id"], checkpoint, exception)
+        if webhook:
+            # TODO add error, values to webhook payload
+            # TODO add retries for webhook calls
+            try:
+                await get_http_client().post(
+                    webhook, json={**run, "status": status}, total_timeout=5
+                )
+            except Exception as e:
+                logger.warning("Failed to send webhook", exc_info=e)
+        # Note we don't handle asyncio.CancelledError here, as we want to
+        # let it bubble up and rollback db transaction, thus marking the run
+        # as available to be picked up by another worker

langgraph_api/route.py

@@ -0,0 +1,118 @@
+import functools
+import inspect
+import typing
+
+import jsonschema_rs
+import orjson
+from starlette._exception_handler import wrap_app_handling_exceptions
+from starlette._utils import is_async_callable
+from starlette.concurrency import run_in_threadpool
+from starlette.middleware import Middleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+from starlette.routing import Route, compile_path, get_name
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from langgraph_api.serde import json_dumpb
+
+
+def api_request_response(
+    func: typing.Callable[[Request], typing.Awaitable[Response] | Response],
+) -> ASGIApp:
+    """
+    Takes a function or coroutine `func(request) -> response`,
+    and returns an ASGI application.
+    """
+
+    async def app(scope: Scope, receive: Receive, send: Send) -> None:
+        request = ApiRequest(scope, receive, send)
+
+        async def app(scope: Scope, receive: Receive, send: Send) -> None:
+            if is_async_callable(func):
+                response = await func(request)
+            else:
+                response = await run_in_threadpool(func, request)
+            await response(scope, receive, send)
+
+        await wrap_app_handling_exceptions(app, request)(scope, receive, send)
+
+    return app
+
+
+class ApiResponse(JSONResponse):
+    def render(self, content: typing.Any) -> bytes:
+        return json_dumpb(content)
+
+
+def _json_loads(
+    content: bytearray, schema: jsonschema_rs.Draft4Validator | None
+) -> typing.Any:
+    json = orjson.loads(content)
+    if schema is not None:
+        schema.validate(json)
+    return json
+
+
+class ApiRequest(Request):
+    async def body(self) -> typing.Coroutine[typing.Any, typing.Any, bytearray]:
+        if not hasattr(self, "_body"):
+            chunks = bytearray()
+            async for chunk in self.stream():
+                chunks.extend(chunk)
+            self._body = chunks
+        return self._body
+
+    async def json(self, schema: jsonschema_rs.Draft4Validator | None) -> typing.Any:
+        if not hasattr(self, "_json"):
+            body = await self.body()
+            self._json = await run_in_threadpool(_json_loads, body, schema)
+        return self._json
+
+
+# ApiRoute uses our custom ApiRequest class to handle requests.
+class ApiRoute(Route):
+    def __init__(
+        self,
+        path: str,
+        endpoint: typing.Callable[..., typing.Any],
+        *,
+        methods: list[str] | None = None,
+        name: str | None = None,
+        include_in_schema: bool = True,
+        middleware: typing.Sequence[Middleware] | None = None,
+    ) -> None:
+        assert path.startswith("/"), "Routed paths must start with '/'"
+        self.path = path
+        self.endpoint = endpoint
+        self.name = get_name(endpoint) if name is None else name
+        self.include_in_schema = include_in_schema
+
+        endpoint_handler = endpoint
+        while isinstance(endpoint_handler, functools.partial):
+            endpoint_handler = endpoint_handler.func
+        if inspect.isfunction(endpoint_handler) or inspect.ismethod(endpoint_handler):
+            # Endpoint is function or method. Treat it as `func(request) -> response`.
+            self.app = api_request_response(endpoint)
+            if methods is None:
+                methods = ["GET"]
+        else:
+            # Endpoint is a class. Treat it as ASGI.
+            self.app = endpoint
+
+        if middleware is not None:
+            for cls, args, kwargs in reversed(middleware):
+                self.app = cls(app=self.app, *args, **kwargs)  # noqa: B026
+
+        if methods is None:
+            self.methods = None
+        else:
+            self.methods = {method.upper() for method in methods}
+            if "GET" in self.methods:
+                self.methods.add("HEAD")
+
+        self.path_regex, self.path_format, self.param_convertors = compile_path(path)
+
+    async def handle(self, scope: Scope, receive: Receive, send: Send) -> None:
+        # https://asgi.readthedocs.io/en/latest/specs/www.html#http-connection-scope
+        scope["route"] = self.path
+        return await super().handle(scope, receive, send)

langgraph_api/schema.py

@@ -0,0 +1,190 @@
+from collections.abc import Sequence
+from datetime import datetime
+from typing import Any, Literal, Optional, TypedDict
+from uuid import UUID
+
+from langgraph_api.serde import Fragment
+
+MetadataInput = dict[str, Any] | None
+MetadataValue = dict[str, Any]
+
+RunStatus = Literal["pending", "error", "success", "timeout", "interrupted"]
+
+ThreadStatus = Literal["idle", "busy", "interrupted", "error"]
+
+StreamMode = Literal["values", "messages", "updates", "events", "debug", "custom"]
+
+MultitaskStrategy = Literal["reject", "rollback", "interrupt", "enqueue"]
+
+OnConflictBehavior = Literal["raise", "do_nothing"]
+
+OnCompletion = Literal["delete", "keep"]
+
+IfNotExists = Literal["create", "reject"]
+
+All = Literal["*"]
+
+
+class Config(TypedDict, total=False):
+    tags: list[str]
+    """
+    Tags for this call and any sub-calls (eg. a Chain calling an LLM).
+    You can use these to filter calls.
+    """
+
+    recursion_limit: int
+    """
+    Maximum number of times a call can recurse. If not provided, defaults to 25.
+    """
+
+    configurable: dict[str, Any]
+    """
+    Runtime values for attributes previously made configurable on this Runnable,
+    or sub-Runnables, through .configurable_fields() or .configurable_alternatives().
+    Check .output_schema() for a description of the attributes that have been made 
+    configurable.
+    """
+
+
+class Checkpoint(TypedDict):
+    thread_id: str
+    checkpoint_ns: str
+    checkpoint_id: str | None
+    checkpoint_map: dict[str, Any] | None
+
+
+class GraphSchema(TypedDict):
+    """Graph model."""
+
+    graph_id: str
+    """The ID of the graph."""
+    state_schema: dict
+    """The schema for the graph state."""
+    config_schema: dict
+    """The schema for the graph config."""
+
+
+class Assistant(TypedDict):
+    """Assistant model."""
+
+    assistant_id: UUID
+    """The ID of the assistant."""
+    graph_id: str
+    """The ID of the graph."""
+    config: Config
+    """The assistant config."""
+    created_at: datetime
+    """The time the assistant was created."""
+    updated_at: datetime
+    """The last time the assistant was updated."""
+    metadata: Fragment
+    """The assistant metadata."""
+    version: int
+    """The assistant version."""
+
+
+class Thread(TypedDict):
+    thread_id: UUID
+    """The ID of the thread."""
+    created_at: datetime
+    """The time the thread was created."""
+    updated_at: datetime
+    """The last time the thread was updated."""
+    metadata: Fragment
+    """The thread metadata."""
+    config: Fragment
+    """The thread config."""
+    status: ThreadStatus
+    """The status of the thread. One of 'idle', 'busy', 'interrupted', "error"."""
+    values: Fragment
+    """The current state of the thread."""
+
+
+class ThreadTask(TypedDict):
+    id: str
+    name: str
+    error: str | None
+    interrupts: list[dict]
+    checkpoint: Checkpoint | None
+    state: Optional["ThreadState"]
+
+
+class ThreadState(TypedDict):
+    values: dict[str, Any]
+    """The state values."""
+    next: Sequence[str]
+    """The name of the node to execute in each task for this step."""
+    checkpoint: Checkpoint
+    """The checkpoint keys. This object can be passed to the /threads and /runs 
+    endpoints to resume execution or update state."""
+    metadata: Fragment
+    """Metadata for this state"""
+    created_at: str | None
+    """Timestamp of state creation"""
+    parent_checkpoint: Checkpoint | None
+    """The parent checkpoint. If missing, this is the root checkpoint."""
+    tasks: Sequence[ThreadTask]
+    """Tasks to execute in this step. If already attempted, may contain an error."""
+
+
+class Run(TypedDict):
+    run_id: UUID
+    """The ID of the run."""
+    thread_id: UUID
+    """The ID of the thread."""
+    assistant_id: UUID
+    """The assistant that was used for this run."""
+    created_at: datetime
+    """The time the run was created."""
+    updated_at: datetime
+    """The last time the run was updated."""
+    status: RunStatus
+    """The status of the run. One of 'pending', 'error', 'success'."""
+    metadata: Fragment
+    """The run metadata."""
+    kwargs: Fragment
+    """The run kwargs."""
+    multitask_strategy: MultitaskStrategy
+    """Strategy to handle concurrent runs on the same thread."""
+
+
+class RunSend(TypedDict):
+    node: str
+    input: dict[str, Any] | None
+
+
+class RunCommand(TypedDict):
+    send: RunSend | Sequence[RunSend] | None
+    update: dict[str, Any] | None
+    resume: Any | None
+
+
+class Cron(TypedDict):
+    """Cron model."""
+
+    cron_id: UUID
+    """The ID of the cron."""
+    thread_id: UUID | None
+    """The ID of the thread."""
+    end_time: datetime | None
+    """The end date to stop running the cron."""
+    schedule: str
+    """The schedule to run, cron format."""
+    created_at: datetime
+    """The time the cron was created."""
+    updated_at: datetime
+    """The last time the cron was updated."""
+    payload: Fragment
+    """The run payload to use for creating new run."""
+
+
+class ThreadUpdateResponse(TypedDict):
+    """Response for updating a thread."""
+
+    checkpoint: Checkpoint
+
+
+class QueueStats(TypedDict):
+    n_pending: int
+    max_age_secs: datetime | None
+    med_age_secs: datetime | None

langgraph_api/serde.py

@@ -0,0 +1,124 @@
+import asyncio
+import pickle
+import uuid
+from base64 import b64encode
+from collections import deque
+from datetime import timedelta, timezone
+from decimal import Decimal
+from ipaddress import (
+    IPv4Address,
+    IPv4Interface,
+    IPv4Network,
+    IPv6Address,
+    IPv6Interface,
+    IPv6Network,
+)
+from pathlib import Path
+from re import Pattern
+from typing import Any, NamedTuple
+from zoneinfo import ZoneInfo
+
+import orjson
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+
+class Fragment(NamedTuple):
+    buf: bytes
+
+
+def decimal_encoder(dec_value: Decimal) -> int | float:
+    """
+    Encodes a Decimal as int of there's no exponent, otherwise float
+
+    This is useful when we use ConstrainedDecimal to represent Numeric(x,0)
+    where a integer (but not int typed) is used. Encoding this as a float
+    results in failed round-tripping between encode and parse.
+    Our Id type is a prime example of this.
+
+    >>> decimal_encoder(Decimal("1.0"))
+    1.0
+
+    >>> decimal_encoder(Decimal("1"))
+    1
+    """
+    if dec_value.as_tuple().exponent >= 0:
+        return int(dec_value)
+    else:
+        return float(dec_value)
+
+
+def default(obj):
+    # Only need to handle types that orjson doesn't serialize by default
+    # https://github.com/ijl/orjson#serialize
+    if isinstance(obj, Fragment):
+        return orjson.Fragment(obj.buf)
+    if hasattr(obj, "model_dump") and callable(obj.model_dump):
+        return obj.model_dump()
+    elif hasattr(obj, "dict") and callable(obj.dict):
+        return obj.dict()
+    elif hasattr(obj, "_asdict") and callable(obj._asdict):
+        return obj._asdict()
+    elif isinstance(obj, BaseException):
+        return {"error": type(obj).__name__, "message": str(obj)}
+    elif isinstance(obj, (set, frozenset, deque)):  # noqa: UP038
+        return list(obj)
+    elif isinstance(obj, (timezone, ZoneInfo)):  # noqa: UP038
+        return obj.tzname(None)
+    elif isinstance(obj, timedelta):
+        return obj.total_seconds()
+    elif isinstance(obj, Decimal):
+        return decimal_encoder(obj)
+    elif isinstance(obj, uuid.UUID):
+        return str(obj)
+    elif isinstance(  # noqa: UP038
+        obj,
+        (
+            IPv4Address,
+            IPv4Interface,
+            IPv4Network,
+            IPv6Address,
+            IPv6Interface,
+            IPv6Network,
+            Path,
+        ),
+    ):
+        return str(obj)
+    elif isinstance(obj, Pattern):
+        return obj.pattern
+    elif isinstance(obj, bytes | bytearray):
+        return b64encode(obj).decode()
+    return None
+
+
+_option = orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS
+
+
+def json_dumpb(obj) -> bytes:
+    return orjson.dumps(obj, default=default, option=_option).replace(
+        b"\u0000", b""
+    )  # null unicode char not allowed in json
+
+
+def json_loads(content: bytes | Fragment | dict) -> Any:
+    if isinstance(content, Fragment):
+        content = content.buf
+    if isinstance(content, dict):
+        return content
+    return orjson.loads(content)
+
+
+async def ajson_loads(content: bytes | Fragment) -> Any:
+    return await asyncio.to_thread(json_loads, content)
+
+
+class Serializer(JsonPlusSerializer):
+    def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
+        try:
+            return super().dumps_typed(obj)
+        except TypeError:
+            return "pickle", pickle.dumps(obj)
+
+    def loads_typed(self, data: tuple[str, bytes]) -> Any:
+        if data[0] == "pickle":
+            return pickle.loads(data[1])
+        return super().loads_typed(data)