langgraph-api 0.4.1__py3-none-any.whl → 0.7.3__py3-none-any.whl

langgraph_api/server.py

@@ -1,38 +1,44 @@
 # MONKEY PATCH: Patch Starlette to fix an error in the library
-# ruff: noqa: E402
-import langgraph_api.patch  # noqa: F401,I001
-import sys
-import os
-
 # WARNING: Keep the import above before other code runs as it
 # patches an error in the Starlette library.
+import langgraph_api.patch  # noqa: F401,I001
+import langgraph_api.timing as timing
 import logging
+import os
+import sys
 import typing
 
 if not (
     (disable_truststore := os.getenv("DISABLE_TRUSTSTORE"))
     and disable_truststore.lower() == "true"
 ):
-    import truststore  # noqa: F401
+    import truststore
 
-    truststore.inject_into_ssl()  # noqa: F401
+    truststore.inject_into_ssl()
 
-from contextlib import asynccontextmanager
 
 import jsonschema_rs
 import structlog
 from langgraph.errors import EmptyInputError, InvalidUpdateError
 from langgraph_sdk.client import configure_loopback_transports
 from starlette.applications import Starlette
+from starlette.exceptions import HTTPException
 from starlette.middleware import Middleware
 from starlette.middleware.cors import CORSMiddleware
-from starlette.routing import Mount
+from starlette.routing import BaseRoute, Mount
 from starlette.types import Receive, Scope, Send
 
 import langgraph_api.config as config
-from langgraph_api.api import meta_routes, routes, user_router
+from langgraph_api.api import (
+    middleware_for_protected_routes,
+    protected_routes,
+    shadowable_meta_routes,
+    unshadowable_meta_routes,
+    user_router,
+)
 from langgraph_api.api.openapi import set_custom_spec
 from langgraph_api.errors import (
+    http_exception_handler,
     overloaded_error_handler,
     validation_error_handler,
     value_error_handler,
@@ -48,21 +54,23 @@ from langgraph_runtime.retry import OVERLOADED_EXCEPTIONS
 logging.captureWarnings(True)
 logger = structlog.stdlib.get_logger(__name__)
 
-middleware = []
+global_middleware = []
 
 if config.ALLOW_PRIVATE_NETWORK:
-    middleware.append(Middleware(PrivateNetworkMiddleware))
+    global_middleware.append(Middleware(PrivateNetworkMiddleware))
 
-if (
+JS_PROXY_MIDDLEWARE_ENABLED = (
     config.HTTP_CONFIG
     and (app := config.HTTP_CONFIG.get("app"))
     and is_js_path(app.split(":")[0])
-):
+)
+
+if JS_PROXY_MIDDLEWARE_ENABLED:
     from langgraph_api.js.remote import JSCustomHTTPProxyMiddleware
 
-    middleware.append(Middleware(JSCustomHTTPProxyMiddleware))
+    global_middleware.append(Middleware(JSCustomHTTPProxyMiddleware))
 
-middleware.extend(
+global_middleware.extend(
     [
         (
             Middleware(
@@ -88,6 +96,7 @@ middleware.extend(
     ]
 )
 exception_handlers = {
+    HTTPException: http_exception_handler,
     ValueError: value_error_handler,
     InvalidUpdateError: value_error_handler,
     EmptyInputError: value_error_handler,
@@ -109,7 +118,7 @@ def update_openapi_spec(app):
         schemas = SchemaGenerator(
             {
                 "openapi": "3.1.0",
-                "info": {"title": "LangGraph Platform", "version": "0.1.0"},
+                "info": {"title": "LangSmith Deployment", "version": "0.1.0"},
             }
         )
         spec = schemas.get_schema(routes=app.routes)
@@ -118,70 +127,135 @@ def update_openapi_spec(app):
         set_custom_spec(spec)
 
 
+def apply_middleware(
+    routes: list[BaseRoute], middleware: list[Middleware]
+) -> list[BaseRoute]:
+    """Applies middleware to a list of routes.
+
+    Routes are modified in place (only the `app` attribute is modified);
+    the modified routes are returned for convenience.
+    """
+    middleware_routes = []
+    for route in routes:
+        for cls, args, kwargs in reversed(middleware):
+            if hasattr(route, "app"):
+                route.app = cls(route.app, *args, **kwargs)  # type: ignore
+            else:
+                raise ValueError(f"Cannot apply middleware: route {route} has no app")
+        middleware_routes.append(route)
+    return middleware_routes
+
+
+custom_middleware = (
+    user_router.user_middleware if user_router and user_router.user_middleware else []
+)
+auth_before_custom_middleware = (
+    config.HTTP_CONFIG and config.HTTP_CONFIG.get("middleware_order") == "auth_first"
+)
+enable_auth_on_custom_routes = config.HTTP_CONFIG and config.HTTP_CONFIG.get(
+    "enable_custom_route_auth"
+)
+# Custom middleware to be applied at the route/mount level, not globally (app level).
+route_level_custom_middleware = (
+    custom_middleware if auth_before_custom_middleware else []
+)
+
+protected_mount = Mount(
+    "",
+    routes=protected_routes,
+    middleware=(
+        middleware_for_protected_routes + route_level_custom_middleware
+        if auth_before_custom_middleware
+        else route_level_custom_middleware + middleware_for_protected_routes
+    ),
+)
+
+
 if user_router:
     # Merge routes
     app = user_router
+    if auth_before_custom_middleware:
+        # Authentication middleware is only applied to protected routes--
+        # it is *not* global middleware. This means that by default,
+        # authentication middleware is necessarily applied *after* any global middleware.
+        # including custom middleware that the user might have supplied.
+        #
+        # To apply authentication middleware before custom middleware,
+        # we must rearrange things a bit:
+        #   1. Extract user-supplied routes and bundle them into a `Mount`
+        #      so that we can easily apply custom middleware to all of them at once.
+        #   2. Extract custom middleware from the user-supplied app.
+        #      Remove it globally, but apply it to each bundle of routes at the mount level.
+        #      This gives us more flexibility in ordering: we can now apply this
+        #      custom middleware before *or* after authentication middleware,
+        #      depending on the `middleware_order` config.
+        user_app = apply_middleware(
+            routes=app.routes,
+            middleware=(
+                middleware_for_protected_routes if enable_auth_on_custom_routes else []
+            )
+            + route_level_custom_middleware,
+        )
+        app.user_middleware = global_middleware
+    else:
+        user_app = (
+            apply_middleware(
+                routes=app.routes,
+                middleware=middleware_for_protected_routes,
+            )
+            if enable_auth_on_custom_routes
+            else app.routes
+        )
+        app.user_middleware = custom_middleware + global_middleware
 
-    meta_route_paths = [
-        getattr(route, "path", None) for route in meta_routes if hasattr(route, "path")
-    ]
-    custom_route_paths = [
-        route.path
-        for route in user_router.router.routes
-        if hasattr(route, "path") and route.path not in meta_route_paths
-    ]
-    logger.info(f"Custom route paths: {custom_route_paths}")
+    app.router.routes = (
+        apply_middleware(unshadowable_meta_routes, route_level_custom_middleware)
+        + user_app
+        + apply_middleware(shadowable_meta_routes, route_level_custom_middleware)
+        + [protected_mount]
+    )
 
     update_openapi_spec(app)
-    for route in routes:
-        if getattr(route, "path", None) in ("/docs", "/openapi.json"):
-            # Our handlers for these are inclusive of the custom routes and default API ones
-            # Don't let these be shadowed
-            app.router.routes.insert(0, route)
-        else:
-            # Everything else could be shadowed.
-            app.router.routes.append(route)
 
-    # Merge lifespans
-    original_lifespan = app.router.lifespan_context
+    # Merge lifespans (base + user)
+    user_lifespan = app.router.lifespan_context
     if app.router.on_startup or app.router.on_shutdown:
         raise ValueError(
             f"Cannot merge lifespans with on_startup or on_shutdown: {app.router.on_startup} {app.router.on_shutdown}"
         )
 
-    @asynccontextmanager
-    async def combined_lifespan(app):
-        async with lifespan(app):
-            if original_lifespan:
-                async with original_lifespan(app):
-                    yield
-            else:
-                yield
-
-    app.router.lifespan_context = combined_lifespan
+    app.router.lifespan_context = timing.combine_lifespans(lifespan, user_lifespan)
 
-    # Merge middleware
-    app.user_middleware = (app.user_middleware or []) + middleware
-    # Merge exception handlers
+    # Merge exception handlers (base + user)
     for k, v in exception_handlers.items():
         if k not in app.exception_handlers:
             app.exception_handlers[k] = v
         else:
             logger.debug(f"Overriding exception handler for {k}")
-    # If the user creates a loopback client with `get_client() (no url)
-    # this will update the http transport to connect to the right app
-    configure_loopback_transports(app)
-
 else:
     # It's a regular starlette app
     app = Starlette(
-        routes=routes,
-        lifespan=lifespan,
-        middleware=middleware,
+        routes=[
+            *apply_middleware(
+                unshadowable_meta_routes + shadowable_meta_routes,
+                route_level_custom_middleware,
+            ),
+            protected_mount,
+        ],
+        lifespan=timing.combine_lifespans(lifespan),
+        middleware=global_middleware,
         exception_handlers=exception_handlers,
     )
 
+# If the user creates a loopback client with `get_client() (no url)
+# this will update the http transport to connect to the right app
+configure_loopback_transports(app)
+
 if config.MOUNT_PREFIX:
+    from starlette.routing import Route
+
+    from langgraph_api.api import meta_metrics, ok
+
     prefix = config.MOUNT_PREFIX
     if not prefix.startswith("/") or prefix.endswith("/"):
         raise ValueError(
@@ -189,8 +263,6 @@ if config.MOUNT_PREFIX:
             f"Valid examples: '/my-api', '/v1', '/api/v1'.\nInvalid examples: 'api/', '/api/'"
         )
     logger.info(f"Mounting routes at prefix: {prefix}")
-    plen = len(prefix)
-    rplen = len(prefix.encode("utf-8"))
 
     class ASGIBypassMiddleware:
         def __init__(self, app: typing.Any, **kwargs):
@@ -208,9 +280,15 @@ if config.MOUNT_PREFIX:
 
             return await self.app(scope, receive, send)
 
+    # Add health checks at root still to avoid having to override health checks.
     app = Starlette(
-        routes=[Mount(prefix, app=app)],
+        routes=[
+            Route("/", ok, methods=["GET"]),
+            Route("/ok", ok, methods=["GET"]),
+            Route("/metrics", meta_metrics, methods=["GET"]),
+            Mount(prefix, app=app),
+        ],
         lifespan=app.router.lifespan_context,
-        middleware=[Middleware(ASGIBypassMiddleware)] + app.user_middleware,
+        middleware=[Middleware(ASGIBypassMiddleware)],
         exception_handlers=app.exception_handlers,
     )

langgraph_api/state.py

@@ -6,12 +6,13 @@ from langgraph.types import Interrupt, StateSnapshot
 
 from langgraph_api.feature_flags import USE_NEW_INTERRUPTS
 from langgraph_api.js.base import RemoteInterrupt
-from langgraph_api.schema import Checkpoint, DeprecatedInterrupt, ThreadState
-from langgraph_api.schema import Interrupt as InterruptSchema
 
 if typing.TYPE_CHECKING:
     from langchain_core.runnables.config import RunnableConfig
 
+    from langgraph_api.schema import Checkpoint, DeprecatedInterrupt, ThreadState
+    from langgraph_api.schema import Interrupt as InterruptSchema
+
 
 def runnable_config_to_checkpoint(
     config: RunnableConfig | None,
@@ -27,7 +28,7 @@ def runnable_config_to_checkpoint(
         return None
 
     configurable = config["configurable"]
-    checkpoint: Checkpoint = {
+    checkpoint: Checkpoint = {  # type: ignore[typed-dict-item]
         "checkpoint_id": configurable["checkpoint_id"],
         "thread_id": configurable["thread_id"],
     }

langgraph_api/store.py

@@ -12,7 +12,8 @@ from langgraph.graph import StateGraph
 from langgraph.pregel import Pregel
 from langgraph.store.base import BaseStore
 
-from langgraph_api import config
+from langgraph_api import config, timing
+from langgraph_api.timing import profiled_import
 from langgraph_api.utils.config import run_in_executor
 
 logger = structlog.stdlib.get_logger(__name__)
@@ -83,22 +84,30 @@ async def collect_store_from_env() -> None:
     CUSTOM_STORE = value
 
 
+@timing.timer(
+    message="Loading store {store_path}",
+    metadata_fn=lambda store_path: {"store_path": store_path},
+    warn_threshold_secs=5,
+    warn_message="Loading store '{store_path}' took longer than expected",
+    error_threshold_secs=10,
+)
 def _load_store(store_path: str) -> Any:
-    if "/" in store_path or ".py:" in store_path:
-        modname = "".join(choice("abcdefghijklmnopqrstuvwxyz") for _ in range(24))
-        path_name, function = store_path.rsplit(":", 1)
-        module_name = path_name.rstrip(":")
-        # Load from file path
-        modspec = importlib.util.spec_from_file_location(modname, module_name)
-        if modspec is None:
-            raise ValueError(f"Could not find store file: {path_name}")
-        module = importlib.util.module_from_spec(modspec)
-        sys.modules[module_name] = module
-        modspec.loader.exec_module(module)  # type: ignore[possibly-unbound-attribute]
-
-    else:
-        path_name, function = store_path.rsplit(".", 1)
-        module = importlib.import_module(path_name)
+    with profiled_import(store_path):
+        if "/" in store_path or ".py:" in store_path:
+            modname = "".join(choice("abcdefghijklmnopqrstuvwxyz") for _ in range(24))
+            path_name, function = store_path.rsplit(":", 1)
+            module_name = path_name.rstrip(":")
+            # Load from file path
+            modspec = importlib.util.spec_from_file_location(modname, module_name)
+            if modspec is None:
+                raise ValueError(f"Could not find store file: {path_name}")
+            module = importlib.util.module_from_spec(modspec)
+            sys.modules[module_name] = module
+            modspec.loader.exec_module(module)  # type: ignore[possibly-unbound-attribute]
+
+        else:
+            path_name, function = store_path.rsplit(".", 1)
+            module = importlib.import_module(path_name)
 
     try:
         store: BaseStore | Callable[[config.StoreConfig], BaseStore] = module.__dict__[

langgraph_api/stream.py

@@ -2,19 +2,20 @@ import uuid
 from collections.abc import AsyncIterator, Callable
 from contextlib import AsyncExitStack, aclosing, asynccontextmanager
 from functools import lru_cache
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import langgraph.version
 import langsmith
 import structlog
 from langchain_core.messages import (
+    AIMessageChunk,
     # TODO: Remove explicit dependency
     BaseMessage,
     BaseMessageChunk,
+    ToolMessageChunk,
     convert_to_messages,
     message_chunk_to_message,
 )
-from langchain_core.runnables import RunnableConfig
 from langgraph.errors import (
     EmptyChannelError,
     EmptyInputError,
@@ -30,7 +31,11 @@ from langgraph_api import __version__
 from langgraph_api import store as api_store
 from langgraph_api.asyncio import ValueEvent, wait_if_not_done
 from langgraph_api.command import map_cmd
-from langgraph_api.feature_flags import USE_RUNTIME_CONTEXT_API
+from langgraph_api.feature_flags import (
+    UPDATES_NEEDED_FOR_INTERRUPTS,
+    USE_DURABILITY,
+    USE_RUNTIME_CONTEXT_API,
+)
 from langgraph_api.graph import get_graph
 from langgraph_api.js.base import BaseRemotePregel
 from langgraph_api.metadata import HOST, PLAN, USER_API_URL, incr_nodes
@@ -40,6 +45,10 @@ from langgraph_api.utils.config import run_in_executor
 from langgraph_runtime.checkpoint import Checkpointer
 from langgraph_runtime.ops import Runs
 
+if TYPE_CHECKING:
+    from langchain_core.runnables import RunnableConfig
+
+
 logger = structlog.stdlib.get_logger(__name__)
 
 
@@ -134,10 +143,18 @@ async def astream_state(
     kwargs = run["kwargs"].copy()
     kwargs.pop("webhook", None)
     kwargs.pop("resumable", False)
+    if USE_DURABILITY:
+        checkpoint_during = kwargs.pop("checkpoint_during")
+        if not kwargs.get("durability") and checkpoint_during:
+            kwargs["durability"] = "async" if checkpoint_during else "exit"
+    else:
+        durability = kwargs.pop("durability")
+        if not kwargs.get("checkpoint_during") and durability in ("async", "exit"):
+            kwargs["checkpoint_during"] = durability == "async"
     subgraphs = kwargs.get("subgraphs", False)
     temporary = kwargs.pop("temporary", False)
     context = kwargs.pop("context", None)
-    config = cast(RunnableConfig, kwargs.pop("config"))
+    config = cast("RunnableConfig", kwargs.pop("config"))
     configurable = config["configurable"]
     stack = AsyncExitStack()
     graph = await stack.enter_async_context(
@@ -146,6 +163,7 @@ async def astream_state(
             config,
             store=(await api_store.get_store()),
             checkpointer=None if temporary else Checkpointer(),
+            is_for_execution=True,
         )
     )
 
@@ -170,7 +188,7 @@ async def astream_state(
     if "messages-tuple" in stream_modes_set and not isinstance(graph, BaseRemotePregel):
         stream_modes_set.remove("messages-tuple")
         stream_modes_set.add("messages")
-    if "updates" not in stream_modes_set:
+    if "updates" not in stream_modes_set and UPDATES_NEEDED_FOR_INTERRUPTS:
         stream_modes_set.add("updates")
         only_interrupt_updates = True
     else:
@@ -221,6 +239,8 @@ async def astream_state(
 
     # stream run
     if use_astream_events:
+        if USE_RUNTIME_CONTEXT_API:
+            kwargs["context"] = context
         async with (
             stack,
             aclosing(  # type: ignore[invalid-argument-type]
@@ -238,7 +258,7 @@ async def astream_state(
                 event = await wait_if_not_done(anext(stream, sentinel), done)
                 if event is sentinel:
                     break
-                event = cast(dict, event)
+                event = cast("dict", event)
                 if event.get("tags") and "langsmith:hidden" in event["tags"]:
                     continue
                 if (
@@ -276,13 +296,25 @@ async def astream_state(
                                 yield "messages", chunk
                         else:
                             msg_, meta = cast(
-                                tuple[BaseMessage | dict, dict[str, Any]], chunk
-                            )
-                            msg = (
-                                convert_to_messages([msg_])[0]
-                                if isinstance(msg_, dict)
-                                else cast(BaseMessage, msg_)
+                                "tuple[BaseMessage | dict, dict[str, Any]]", chunk
                             )
+                            is_chunk = False
+                            if isinstance(msg_, dict):
+                                if (
+                                    "chunk" in msg_.get("type", "").lower()
+                                    or "chunk" in msg_.get("role", "").lower()
+                                ):
+                                    if "ai" in msg_.get("role", "").lower():
+                                        msg = AIMessageChunk(**msg_)  # type: ignore[arg-type]
+                                    elif "tool" in msg_.get("role", "").lower():
+                                        msg = ToolMessageChunk(**msg_)  # type: ignore[arg-type]
+                                    else:
+                                        msg = BaseMessageChunk(**msg_)  # type: ignore[arg-type]
+                                    is_chunk = True
+                                else:
+                                    msg = convert_to_messages([msg_])[0]
+                            else:
+                                msg = msg_
                             if msg.id in messages:
                                 messages[msg.id] += msg
                             else:
@@ -294,7 +326,13 @@ async def astream_state(
                                     if isinstance(msg, BaseMessageChunk)
                                     else "messages/complete"
                                 ),
-                                [message_chunk_to_message(messages[msg.id])],
+                                [
+                                    (
+                                        message_chunk_to_message(messages[msg.id])
+                                        if not is_chunk
+                                        else messages[msg.id]
+                                    )
+                                ],
                             )
                     elif mode in stream_mode:
                         if subgraphs and ns:
@@ -308,11 +346,9 @@ async def astream_state(
                         and len(chunk["__interrupt__"]) > 0
                         and only_interrupt_updates
                     ):
-                        # We always want to return interrupt events by default.
-                        # If updates aren't specified as a stream mode, we return these as values events.
                         # If the interrupt doesn't have any actions (e.g. interrupt before or after a node is specified), we don't return the interrupt at all today.
                         if subgraphs and ns:
-                            yield f"values|{'|'.join(ns)}", chunk
+                            yield "values|{'|'.join(ns)}", chunk
                         else:
                             yield "values", chunk
                     # --- end shared logic with astream ---
@@ -340,9 +376,9 @@ async def astream_state(
                 if event is sentinel:
                     break
                 if subgraphs:
-                    ns, mode, chunk = cast(tuple[str, str, dict[str, Any]], event)
+                    ns, mode, chunk = cast("tuple[str, str, dict[str, Any]]", event)
                 else:
-                    mode, chunk = cast(tuple[str, dict[str, Any]], event)
+                    mode, chunk = cast("tuple[str, dict[str, Any]]", event)
                     ns = None
                 # --- begin shared logic with astream_events ---
                 if mode == "debug":
@@ -360,14 +396,25 @@ async def astream_state(
                             yield "messages", chunk
                     else:
                         msg_, meta = cast(
-                            tuple[BaseMessage | dict, dict[str, Any]], chunk
-                        )
-                        msg = (
-                            convert_to_messages([msg_])[0]
-                            if isinstance(msg_, dict)
-                            else cast(BaseMessage, msg_)
+                            "tuple[BaseMessage | dict, dict[str, Any]]", chunk
                         )
-
+                        is_chunk = False
+                        if isinstance(msg_, dict):
+                            if (
+                                "chunk" in msg_.get("type", "").lower()
+                                or "chunk" in msg_.get("role", "").lower()
+                            ):
+                                if "ai" in msg_.get("role", "").lower():
+                                    msg = AIMessageChunk(**msg_)  # type: ignore[arg-type]
+                                elif "tool" in msg_.get("role", "").lower():
+                                    msg = ToolMessageChunk(**msg_)  # type: ignore[arg-type]
+                                else:
+                                    msg = BaseMessageChunk(**msg_)  # type: ignore[arg-type]
+                                is_chunk = True
+                            else:
+                                msg = convert_to_messages([msg_])[0]
+                        else:
+                            msg = msg_
                         if msg.id in messages:
                             messages[msg.id] += msg
                         else:
@@ -379,7 +426,13 @@ async def astream_state(
                                 if isinstance(msg, BaseMessageChunk)
                                 else "messages/complete"
                             ),
-                            [message_chunk_to_message(messages[msg.id])],
+                            [
+                                (
+                                    message_chunk_to_message(messages[msg.id])
+                                    if not is_chunk
+                                    else messages[msg.id]
+                                )
+                            ],
                         )
                 elif mode in stream_mode:
                     if subgraphs and ns:
@@ -393,11 +446,9 @@ async def astream_state(
                     and len(chunk["__interrupt__"]) > 0
                     and only_interrupt_updates
                 ):
-                    # We always want to return interrupt events by default.
-                    # If updates aren't specified as a stream mode, we return these as values events.
                     # If the interrupt doesn't have any actions (e.g. interrupt before or after a node is specified), we don't return the interrupt at all today.
                     if subgraphs and ns:
-                        yield "values|{'|'.join(ns)}", chunk
+                        yield f"values|{'|'.join(ns)}", chunk
                     else:
                         yield "values", chunk
                 # --- end shared logic with astream_events ---