langgraph-api 0.0.1__py3-none-any.whl

langgraph_api/server.py

@@ -0,0 +1,48 @@
+import logging
+
+import jsonschema_rs
+import structlog
+from langgraph.errors import EmptyInputError, InvalidUpdateError
+from starlette.applications import Starlette
+from starlette.middleware import Middleware
+from starlette.middleware.cors import CORSMiddleware
+
+import langgraph_api.config as config
+import langgraph_api.patch  # noqa: F401
+from langgraph_api.api import routes
+from langgraph_api.errors import (
+    overloaded_error_handler,
+    validation_error_handler,
+    value_error_handler,
+)
+from langgraph_api.http_logger import AccessLoggerMiddleware
+from langgraph_api.lifespan import lifespan
+from langgraph_license.middleware import LicenseValidationMiddleware
+from langgraph_storage.retry import OVERLOADED_EXCEPTIONS
+
+logging.captureWarnings(True)
+logger = structlog.stdlib.get_logger(__name__)
+
+
+app = Starlette(
+    routes=routes,
+    lifespan=lifespan,
+    middleware=[
+        Middleware(
+            CORSMiddleware,
+            allow_origins=config.CORS_ALLOW_ORIGINS,
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        ),
+        Middleware(LicenseValidationMiddleware),
+        Middleware(AccessLoggerMiddleware, logger=logger),
+    ],
+    exception_handlers={
+        ValueError: value_error_handler,
+        InvalidUpdateError: value_error_handler,
+        EmptyInputError: value_error_handler,
+        jsonschema_rs.ValidationError: validation_error_handler,
+    }
+    | {exc: overloaded_error_handler for exc in OVERLOADED_EXCEPTIONS},
+)

langgraph_api/sse.py

@@ -0,0 +1,118 @@
+import asyncio
+from collections.abc import AsyncIterator, Awaitable, Callable, Mapping
+from functools import partial
+from typing import Any
+
+import anyio
+import sse_starlette
+import sse_starlette.sse
+import structlog.stdlib
+from starlette.types import Receive, Scope, Send
+
+from langgraph_api.asyncio import SimpleTaskGroup, aclosing
+from langgraph_api.serde import json_dumpb
+
+logger = structlog.stdlib.get_logger(__name__)
+
+
+class EventSourceResponse(sse_starlette.EventSourceResponse):
+    def __init__(
+        self,
+        content: AsyncIterator[bytes | tuple[bytes, Any | bytes]],
+        status_code: int = 200,
+        headers: Mapping[str, str] | None = None,
+    ) -> None:
+        super().__init__(content=content, status_code=status_code, headers=headers)
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        async with anyio.create_task_group() as task_group:
+            # https://trio.readthedocs.io/en/latest/reference-core.html#custom-supervisors
+            async def wrap(func: Callable[[], Awaitable[None]]) -> None:
+                await func()
+                # noinspection PyAsyncCall
+                task_group.cancel_scope.cancel()
+
+            task_group.start_soon(wrap, partial(self.stream_response, send))
+            task_group.start_soon(wrap, self.listen_for_exit_signal)
+
+            if self.data_sender_callable:
+                task_group.start_soon(self.data_sender_callable)
+
+            await wrap(partial(self.listen_for_disconnect, receive))
+
+        if self.background is not None:  # pragma: no cover, tested in StreamResponse
+            await self.background()
+
+    async def stream_response(self, send: Send) -> None:
+        await send(
+            {
+                "type": "http.response.start",
+                "status": self.status_code,
+                "headers": self.raw_headers,
+            }
+        )
+        async with (
+            SimpleTaskGroup(sse_heartbeat(send), cancel=True, wait=False),
+            aclosing(self.body_iterator) as body,
+        ):
+            try:
+                async for data in body:
+                    with anyio.move_on_after(self.send_timeout) as timeout:
+                        await send(
+                            {
+                                "type": "http.response.body",
+                                "body": json_to_sse(*data)
+                                if isinstance(data, tuple)
+                                else data,
+                                "more_body": True,
+                            }
+                        )
+                    if timeout.cancel_called:
+                        raise sse_starlette.sse.SendTimeoutError()
+            except sse_starlette.sse.SendTimeoutError:
+                raise
+            except Exception as exc:
+                await logger.aexception("Error streaming response", exc_info=exc)
+                await send(
+                    {
+                        "type": "http.response.body",
+                        "body": json_to_sse(b"error", exc),
+                        "more_body": False,
+                    }
+                )
+
+        async with self._send_lock:
+            self.active = False
+            await send({"type": "http.response.body", "body": b"", "more_body": False})
+
+
+async def sse_heartbeat(send: Send) -> None:
+    payload = sse_starlette.ServerSentEvent(comment="heartbeat").encode()
+    try:
+        while True:
+            await asyncio.sleep(5)
+            await send(
+                {"type": "http.response.body", "body": payload, "more_body": True}
+            )
+    except asyncio.CancelledError:
+        pass
+
+
+SEP = b"\r\n"
+EVENT = b"event: "
+DATA = b"data: "
+BYTES_LIKE = (bytes, bytearray, memoryview)
+
+
+def json_to_sse(event: bytes, data: Any | bytes) -> bytes:
+    return b"".join(
+        (
+            EVENT,
+            event,
+            SEP,
+            DATA,
+            data if isinstance(data, BYTES_LIKE) else json_dumpb(data),
+            SEP,
+            SEP,
+        )
+    )

langgraph_api/state.py

@@ -0,0 +1,67 @@
+from langchain_core.runnables.config import RunnableConfig
+from langgraph.types import StateSnapshot
+
+from langgraph_api.schema import Checkpoint, ThreadState
+
+
+def runnable_config_to_checkpoint(
+    config: RunnableConfig | None,
+) -> Checkpoint | None:
+    if (
+        not config
+        or not config["configurable"]
+        or "thread_id" not in config["configurable"]
+        or not config["configurable"]["thread_id"]
+        or "checkpoint_id" not in config["configurable"]
+        or not config["configurable"]["checkpoint_id"]
+    ):
+        return None
+
+    configurable = config["configurable"]
+    checkpoint: Checkpoint = {
+        "checkpoint_id": configurable["checkpoint_id"],
+        "thread_id": configurable["thread_id"],
+    }
+
+    if "checkpoint_ns" in configurable:
+        checkpoint["checkpoint_ns"] = configurable["checkpoint_ns"] or ""
+
+    if "checkpoint_map" in configurable:
+        checkpoint["checkpoint_map"] = configurable["checkpoint_map"]
+
+    return checkpoint
+
+
+def state_snapshot_to_thread_state(state: StateSnapshot) -> ThreadState:
+    return {
+        "values": state.values,
+        "next": state.next,
+        "tasks": [
+            {
+                "id": t.id,
+                "name": t.name,
+                "path": t.path,
+                "error": t.error,
+                "interrupts": t.interrupts,
+                "checkpoint": t.state["configurable"]
+                if t.state is not None and not isinstance(t.state, StateSnapshot)
+                else None,
+                "state": state_snapshot_to_thread_state(t.state)
+                if isinstance(t.state, StateSnapshot)
+                else None,
+                "result": getattr(t, "result", None),
+            }
+            for t in state.tasks
+        ],
+        "metadata": state.metadata,
+        "created_at": state.created_at,
+        "checkpoint": runnable_config_to_checkpoint(state.config),
+        "parent_checkpoint": runnable_config_to_checkpoint(state.parent_config),
+        # below are deprecated
+        "checkpoint_id": state.config["configurable"].get("checkpoint_id")
+        if state.config
+        else None,
+        "parent_checkpoint_id": state.parent_config["configurable"]["checkpoint_id"]
+        if state.parent_config
+        else None,
+    }

langgraph_api/stream.py

@@ -0,0 +1,289 @@
+from collections.abc import AsyncIterator, Callable
+from contextlib import AsyncExitStack, aclosing
+from functools import lru_cache
+from typing import Any, cast
+
+import langgraph.version
+import langsmith
+import structlog
+from langchain_core.messages import (
+    BaseMessage,
+    BaseMessageChunk,
+    message_chunk_to_message,
+)
+from langchain_core.runnables.config import run_in_executor
+from langgraph.errors import (
+    EmptyChannelError,
+    EmptyInputError,
+    GraphRecursionError,
+    InvalidUpdateError,
+)
+from langgraph.pregel.debug import CheckpointPayload
+from langgraph.types import Command, Send
+from pydantic import ValidationError
+from pydantic.v1 import ValidationError as ValidationErrorLegacy
+
+from langgraph_api.asyncio import ValueEvent, wait_if_not_done
+from langgraph_api.graph import get_graph
+from langgraph_api.js.remote import RemotePregel
+from langgraph_api.metadata import HOST, PLAN, incr_nodes
+from langgraph_api.schema import Run, RunCommand, StreamMode
+from langgraph_api.serde import json_dumpb
+from langgraph_api.utils import AsyncConnectionProto
+from langgraph_storage.checkpoint import Checkpointer
+from langgraph_storage.ops import Runs
+from langgraph_storage.store import Store
+
+logger = structlog.stdlib.get_logger(__name__)
+
+AnyStream = AsyncIterator[tuple[str, Any]]
+
+
+def _preproces_debug_checkpoint_task(task: dict[str, Any]) -> dict[str, Any]:
+    if (
+        "state" not in task
+        or not task["state"]
+        or "configurable" not in task["state"]
+        or not task["state"]["configurable"]
+    ):
+        return task
+
+    task["checkpoint"] = task["state"]["configurable"]
+    del task["state"]
+    return task
+
+
+def _preprocess_debug_checkpoint(payload: CheckpointPayload | None) -> dict[str, Any]:
+    from langgraph_api.state import runnable_config_to_checkpoint
+
+    if not payload:
+        return None
+
+    payload["checkpoint"] = runnable_config_to_checkpoint(payload["config"])
+    payload["parent_checkpoint"] = runnable_config_to_checkpoint(
+        payload["parent_config"] if "parent_config" in payload else None
+    )
+
+    payload["tasks"] = [_preproces_debug_checkpoint_task(t) for t in payload["tasks"]]
+
+    # TODO: deprecate the `config`` and `parent_config`` fields
+    return payload
+
+
+def _map_cmd(cmd: RunCommand) -> Command:
+    send = cmd.get("send")
+    if send is not None and not isinstance(send, list):
+        send = [cmd.get("send")]
+
+    return Command(
+        update=cmd.get("update"),
+        send=[Send(send["node"], send["input"]) for send in send] if send else None,
+        resume=cmd.get("resume"),
+    )
+
+
+async def astream_state(
+    stack: AsyncExitStack,
+    conn: AsyncConnectionProto,
+    run: Run,
+    attempt: int,
+    done: ValueEvent,
+    *,
+    on_checkpoint: Callable[[CheckpointPayload], None] = lambda _: None,
+) -> AnyStream:
+    """Stream messages from the runnable."""
+    run_id = str(run["run_id"])
+    await stack.enter_async_context(conn.pipeline())
+    # extract args from run
+    kwargs = run["kwargs"].copy()
+    subgraphs = kwargs.get("subgraphs", False)
+    temporary = kwargs.pop("temporary", False)
+    config = kwargs.pop("config")
+    graph = get_graph(
+        config["configurable"]["graph_id"],
+        config,
+        store=None if not conn else Store(conn),
+        checkpointer=None if temporary else Checkpointer(conn),
+    )
+    input = kwargs.pop("input")
+    if cmd := kwargs.pop("command"):
+        input = _map_cmd(cmd)
+    stream_mode: list[StreamMode] = kwargs.pop("stream_mode")
+    feedback_keys = kwargs.pop("feedback_keys", None)
+    stream_modes_set: set[StreamMode] = set(stream_mode) - {"events"}
+    if "debug" not in stream_modes_set:
+        stream_modes_set.add("debug")
+    if "messages-tuple" in stream_modes_set:
+        stream_modes_set.remove("messages-tuple")
+        stream_modes_set.add("messages")
+    # attach attempt metadata
+    config["metadata"]["run_attempt"] = attempt
+    # attach langgraph metadata
+    config["metadata"]["langgraph_version"] = langgraph.version.__version__
+    config["metadata"]["langgraph_plan"] = PLAN
+    config["metadata"]["langgraph_host"] = HOST
+    # attach node counter
+    if not isinstance(graph, RemotePregel):
+        config["configurable"]["__pregel_node_finished"] = incr_nodes
+        # TODO add node tracking for JS graphs
+    # attach run_id to config
+    # for attempts beyond the first, use a fresh, unique run_id
+    config = {**config, "run_id": run["run_id"]} if attempt == 1 else config
+    # set up state
+    checkpoint: CheckpointPayload | None = None
+    messages: dict[str, BaseMessageChunk] = {}
+    use_astream_events = "events" in stream_mode or isinstance(graph, RemotePregel)
+    # yield metadata chunk
+    yield "metadata", {"run_id": run_id, "attempt": attempt}
+    # stream run
+    if use_astream_events:
+        async with (
+            stack,
+            aclosing(
+                graph.astream_events(
+                    input,
+                    config,
+                    version="v2",
+                    stream_mode=list(stream_modes_set),
+                    **kwargs,
+                )
+            ) as stream,
+        ):
+            sentinel = object()
+            while True:
+                event = await wait_if_not_done(anext(stream, sentinel), done)
+                if event is sentinel:
+                    break
+                if event.get("tags") and "langsmith:hidden" in event["tags"]:
+                    continue
+                if "messages" in stream_mode and isinstance(graph, RemotePregel):
+                    if event["event"] == "on_custom_event" and event["name"] in (
+                        "messages/complete",
+                        "messages/partial",
+                        "messages/metadata",
+                    ):
+                        yield event["name"], event["data"]
+                # TODO support messages-tuple for js graphs
+                if event["event"] == "on_chain_stream" and event["run_id"] == run_id:
+                    if subgraphs:
+                        ns, mode, chunk = event["data"]["chunk"]
+                    else:
+                        mode, chunk = event["data"]["chunk"]
+                    # --- begin shared logic with astream ---
+                    if mode == "debug":
+                        if chunk["type"] == "checkpoint":
+                            checkpoint = _preprocess_debug_checkpoint(chunk["payload"])
+                            on_checkpoint(checkpoint)
+                    if mode == "messages":
+                        if "messages-tuple" in stream_mode:
+                            yield "messages", chunk
+                        else:
+                            msg, meta = cast(tuple[BaseMessage, dict[str, Any]], chunk)
+                            if msg.id in messages:
+                                messages[msg.id] += msg
+                            else:
+                                messages[msg.id] = msg
+                                yield "messages/metadata", {msg.id: {"metadata": meta}}
+                            yield (
+                                "messages/partial"
+                                if isinstance(msg, BaseMessageChunk)
+                                else "messages/complete",
+                                [message_chunk_to_message(messages[msg.id])],
+                            )
+                    elif mode in stream_mode:
+                        if subgraphs and ns:
+                            yield f"{mode}|{'|'.join(ns)}", chunk
+                        else:
+                            yield mode, chunk
+                    # --- end shared logic with astream ---
+                elif "events" in stream_mode:
+                    yield "events", event
+    else:
+        async with (
+            stack,
+            aclosing(
+                graph.astream(
+                    input, config, stream_mode=list(stream_modes_set), **kwargs
+                )
+            ) as stream,
+        ):
+            sentinel = object()
+            while True:
+                event = await wait_if_not_done(anext(stream, sentinel), done)
+                if event is sentinel:
+                    break
+                if subgraphs:
+                    ns, mode, chunk = event
+                else:
+                    mode, chunk = event
+                # --- begin shared logic with astream_events ---
+                if mode == "debug":
+                    if chunk["type"] == "checkpoint":
+                        checkpoint = _preprocess_debug_checkpoint(chunk["payload"])
+                        on_checkpoint(checkpoint)
+                if mode == "messages":
+                    if "messages-tuple" in stream_mode:
+                        yield "messages", chunk
+                    else:
+                        msg, meta = cast(tuple[BaseMessage, dict[str, Any]], chunk)
+                        if msg.id in messages:
+                            messages[msg.id] += msg
+                        else:
+                            messages[msg.id] = msg
+                            yield "messages/metadata", {msg.id: {"metadata": meta}}
+                        yield (
+                            "messages/partial"
+                            if isinstance(msg, BaseMessageChunk)
+                            else "messages/complete",
+                            [message_chunk_to_message(messages[msg.id])],
+                        )
+                elif mode in stream_mode:
+                    if subgraphs and ns:
+                        yield f"{mode}|{'|'.join(ns)}", chunk
+                    else:
+                        yield mode, chunk
+                # --- end shared logic with astream_events ---
+    # Get feedback URLs
+    if feedback_keys:
+        feedback_urls = await run_in_executor(
+            None, get_feedback_urls, run_id, feedback_keys
+        )
+        yield "feedback", feedback_urls
+
+
+async def consume(stream: AnyStream, run_id: str) -> None:
+    async with aclosing(stream):
+        try:
+            async for mode, payload in stream:
+                await Runs.Stream.publish(
+                    run_id, mode, await run_in_executor(None, json_dumpb, payload)
+                )
+        except Exception as e:
+            if isinstance(e, ExceptionGroup):
+                e = e.exceptions[0]
+            await Runs.Stream.publish(
+                run_id, "error", await run_in_executor(None, json_dumpb, e)
+            )
+            raise e from None
+
+
+def get_feedback_urls(run_id: str, feedback_keys: list[str]) -> dict[str, str]:
+    client = get_langsmith_client()
+    tokens = client.create_presigned_feedback_tokens(run_id, feedback_keys)
+    return {key: token.url for key, token in zip(feedback_keys, tokens, strict=False)}
+
+
+@lru_cache(maxsize=1)
+def get_langsmith_client() -> langsmith.Client:
+    return langsmith.Client()
+
+
+EXPECTED_ERRORS = (
+    ValueError,
+    InvalidUpdateError,
+    GraphRecursionError,
+    EmptyInputError,
+    EmptyChannelError,
+    ValidationError,
+    ValidationErrorLegacy,
+)

langgraph_api/utils.py

@@ -0,0 +1,60 @@
+import uuid
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from datetime import datetime
+from typing import Any, Protocol, TypeAlias, TypeVar
+
+from starlette.exceptions import HTTPException
+
+T = TypeVar("T")
+Row: TypeAlias = dict[str, Any]
+
+
+class AsyncCursorProto(Protocol):
+    async def fetchone(self) -> Row: ...
+
+    async def fetchall(self) -> list[Row]: ...
+
+    async def __aiter__(self) -> AsyncIterator[Row]:
+        yield ...
+
+
+class AsyncPipelineProto(Protocol):
+    async def sync(self) -> None: ...
+
+
+class AsyncConnectionProto(Protocol):
+    @asynccontextmanager
+    async def pipeline(self) -> AsyncIterator[AsyncPipelineProto]:
+        yield ...
+
+    async def execute(self, query: str, *args, **kwargs) -> AsyncCursorProto: ...
+
+
+async def fetchone(
+    it: AsyncIterator[T],
+    *,
+    not_found_code: int = 404,
+    not_found_detail: str | None = None,
+) -> T:
+    """Fetch the first row from an async iterator."""
+    try:
+        return await anext(it)
+    except StopAsyncIteration:
+        raise HTTPException(
+            status_code=not_found_code, detail=not_found_detail
+        ) from None
+
+
+def validate_uuid(uuid_str: str, invalid_uuid_detail: str | None) -> uuid.UUID:
+    try:
+        return uuid.UUID(uuid_str)
+    except ValueError:
+        raise HTTPException(status_code=422, detail=invalid_uuid_detail) from None
+
+
+def next_cron_date(schedule: str, base_time: datetime) -> datetime:
+    import croniter
+
+    cron_iter = croniter.croniter(schedule, base_time)
+    return cron_iter.get_next(datetime)