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

langgraph_api/timing/timer.py

@@ -0,0 +1,318 @@
+from __future__ import annotations
+
+# ruff: noqa: E402
+import time
+
+# Capture the time when this module is first imported (early in server startup)
+_PROCESS_START_TIME = time.monotonic()
+
+import contextlib
+import functools
+import inspect
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Generic, ParamSpec, TypeVar, overload
+
+import structlog
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from contextlib import AbstractAsyncContextManager
+
+    from starlette.applications import Starlette
+
+logger = structlog.stdlib.get_logger(__name__)
+
+P = ParamSpec("P")
+R = TypeVar("R")
+T = TypeVar("T", covariant=True)
+
+
+@dataclass(frozen=True)
+class TimerConfig(Generic[P]):
+    message: str = "Function timing"
+    metadata_fn: Callable[P, dict[str, Any]] | None = None
+    warn_threshold_secs: float | None = None
+    warn_message: str | None = None
+    error_threshold_secs: float | None = None
+    error_message: str | None = None
+
+
+@overload
+def timer(_func: Callable[P, R], /, **kwargs) -> Callable[P, R]: ...
+@overload
+def timer(
+    _func: None = None, /, **kwargs
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+
+def timer(
+    _func: Callable[P, R] | None = None,
+    /,
+    message: str = "Function timing",
+    metadata_fn: Callable[P, dict[str, Any]] | None = None,
+    warn_threshold_secs: float | None = None,
+    warn_message: str | None = None,
+    error_threshold_secs: float | None = None,
+    error_message: str | None = None,
+):
+    """
+    Decorator for sync *and* async callables.
+    """
+    cfg = TimerConfig[P](
+        message=message,
+        metadata_fn=metadata_fn,
+        warn_threshold_secs=warn_threshold_secs,
+        warn_message=warn_message,
+        error_threshold_secs=error_threshold_secs,
+        error_message=error_message,
+    )
+
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def awrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+                start = time.perf_counter()
+                exc: BaseException | None = None
+                try:
+                    return await func(*args, **kwargs)  # type: ignore[misc]
+                except BaseException as e:
+                    exc = e
+                    raise
+                finally:
+                    elapsed = time.perf_counter() - start
+                    _log_timing(
+                        name=func.__qualname__,
+                        elapsed=elapsed,
+                        cfg=cfg,  # type: ignore[arg-type]
+                        args=args,
+                        kwargs=kwargs,
+                        exc=exc,
+                    )
+
+            return awrapper  # type: ignore[return-value]
+
+        @functools.wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            start = time.perf_counter()
+            exc: BaseException | None = None
+            try:
+                return func(*args, **kwargs)
+            except BaseException as e:
+                exc = e
+                raise
+            finally:
+                elapsed = time.perf_counter() - start
+                _log_timing(
+                    name=func.__qualname__,
+                    elapsed=elapsed,
+                    cfg=cfg,  # type: ignore[arg-type]
+                    args=args,
+                    kwargs=kwargs,
+                    exc=exc,
+                )
+
+        return wrapper
+
+    return decorator(_func) if _func is not None else decorator
+
+
+class aenter_timed(Generic[T]):
+    """
+    Wraps an async context manager and logs the time spent in *its __aenter__*.
+    __aexit__ is delegated without additional timing.
+    """
+
+    def __init__(
+        self,
+        inner: AbstractAsyncContextManager[T],
+        *,
+        name: str,
+        cfg: TimerConfig[Any],
+        extra: dict[str, Any] | None = None,
+    ) -> None:
+        self._inner = inner
+        self._name = name
+        self._cfg = cfg
+        self._extra = extra or {}
+
+    async def __aenter__(self) -> T:
+        start = time.perf_counter()
+        exc: BaseException | None = None
+        try:
+            return await self._inner.__aenter__()
+        except BaseException as e:
+            exc = e
+            raise
+        finally:
+            elapsed = time.perf_counter() - start
+            _log_timing(
+                name=self._name,
+                elapsed=elapsed,
+                cfg=self._cfg,
+                exc=exc,
+                extra=self._extra,
+            )
+
+    async def __aexit__(self, exc_type, exc, tb) -> bool | None:
+        return await self._inner.__aexit__(exc_type, exc, tb)
+
+
+def time_aenter(
+    cm: AbstractAsyncContextManager[T],
+    *,
+    name: str,
+    message: str,
+    warn_threshold_secs: float | None = None,
+    warn_message: str | None = None,
+    error_threshold_secs: float | None = None,
+    error_message: str | None = None,
+    extra: dict[str, Any] | None = None,
+) -> aenter_timed[T]:
+    """
+    Convenience helper to wrap any async CM and time only its __aenter__.
+    """
+    cfg = TimerConfig[Any](
+        message=message,
+        warn_threshold_secs=warn_threshold_secs,
+        warn_message=warn_message,
+        error_threshold_secs=error_threshold_secs,
+        error_message=error_message,
+        metadata_fn=None,
+    )
+    return aenter_timed(cm, name=name, cfg=cfg, extra=extra)
+
+
+def wrap_lifespan_context_aenter(
+    lifespan_ctx: Callable[[Any], AbstractAsyncContextManager[Any]],
+    *,
+    name: str = "user_router.lifespan",
+    message: str = "Entered lifespan context",
+    warn_threshold_secs: float | None = 10,
+    warn_message: str | None = (
+        "User lifespan startup exceeded expected time. "
+        "Slow work done at entry time within lifespan context can delay readiness, "
+        "reduce scale-out capacity, and may cause deployments to be marked unhealthy."
+    ),
+    error_threshold_secs: float | None = 30,
+    error_message: str | None = None,
+) -> Callable[[Any], AbstractAsyncContextManager[Any]]:
+    @functools.wraps(lifespan_ctx)
+    def wrapped(app: Any) -> AbstractAsyncContextManager[Any]:
+        return time_aenter(
+            lifespan_ctx(app),
+            name=name,
+            message=message,
+            warn_threshold_secs=warn_threshold_secs,
+            warn_message=warn_message,
+            error_threshold_secs=error_threshold_secs,
+            error_message=error_message,
+        )
+
+    return wrapped
+
+
+LP = ParamSpec("LP")
+
+
+def combine_lifespans(
+    *lifespans: Callable[[Starlette], AbstractAsyncContextManager] | None,
+) -> Callable[[Starlette], AbstractAsyncContextManager]:
+    @contextlib.asynccontextmanager
+    async def combined_lifespan(app):
+        async with contextlib.AsyncExitStack() as stack:
+            for ls in lifespans:
+                if ls is not None:
+                    await stack.enter_async_context(ls(app))
+            elapsed = get_startup_elapsed()
+            logger.info(f"Application started up in {elapsed:2.3f}s", elapsed=elapsed)
+            yield
+
+    return combined_lifespan
+
+
+def get_startup_elapsed() -> float:
+    """Return elapsed seconds since the process started (module import time)."""
+    return time.monotonic() - _PROCESS_START_TIME
+
+
+def _log_timing(
+    *,
+    name: str,
+    elapsed: float,
+    cfg: TimerConfig[Any],
+    args: tuple[Any, ...] = (),
+    kwargs: dict[str, Any] | None = None,
+    exc: BaseException | None = None,
+    extra: dict[str, Any] | None = None,
+) -> None:
+    from langgraph_api import config
+
+    kwargs = kwargs or {}
+
+    log_data: dict[str, Any] = {
+        "name": name,
+        "elapsed_seconds": elapsed,
+    }
+
+    if extra:
+        log_data.update(extra)
+
+    if cfg.metadata_fn is not None:
+        try:
+            md = cfg.metadata_fn(*args, **kwargs)  # type: ignore[misc]
+            if not isinstance(md, dict):
+                raise TypeError("metadata_fn must return a dict")
+            log_data.update(md)
+        except Exception as meta_exc:
+            log_data["metadata_error"] = repr(meta_exc)
+
+    if exc is not None:
+        log_data["exception"] = repr(exc)
+
+    level, msg = _pick_level_and_message(
+        elapsed=elapsed,
+        message=cfg.message,
+        warn_threshold_secs=cfg.warn_threshold_secs,
+        warn_message=cfg.warn_message,
+        error_threshold_secs=cfg.error_threshold_secs,
+        error_message=cfg.error_message,
+    )
+
+    # Allow {graph_id} etc.
+    msg = msg.format(**log_data)
+
+    # Add profiler hint if we hit warn/error threshold and profiling isn't enabled
+    if level >= logging.WARNING and not config.FF_PROFILE_IMPORTS:
+        msg = (
+            f"{msg}\n"
+            f"    To get detailed profiling of slow operations, set FF_PROFILE_IMPORTS=true"
+        )
+
+    logger.log(level, msg, **log_data)
+
+
+def _pick_level_and_message(
+    *,
+    elapsed: float,
+    message: str,
+    warn_threshold_secs: float | None,
+    warn_message: str | None,
+    error_threshold_secs: float | None,
+    error_message: str | None,
+) -> tuple[int, str]:
+    level = logging.INFO
+    msg = message
+
+    if warn_threshold_secs is not None and elapsed > warn_threshold_secs:
+        level = logging.WARNING
+        if warn_message is not None:
+            msg = warn_message
+
+    if error_threshold_secs is not None and elapsed > error_threshold_secs:
+        level = logging.ERROR
+        if error_message is not None:
+            msg = error_message
+
+    return level, msg

langgraph_api/utils/__init__.py

@@ -7,6 +7,7 @@ from datetime import datetime
 from typing import Any, Protocol, TypeAlias, TypeVar, cast
 
 import structlog
+from langchain_core.runnables import RunnableConfig
 from langgraph_sdk import Auth
 from starlette.authentication import AuthCredentials, BaseUser
 from starlette.exceptions import HTTPException
@@ -36,7 +37,7 @@ async def with_user(
     if current is None:
         return
     set_auth_ctx(
-        cast(BaseUser, current.user), AuthCredentials(scopes=current.permissions)
+        cast("BaseUser", current.user), AuthCredentials(scopes=current.permissions)
     )
 
 
@@ -60,6 +61,42 @@ def get_auth_ctx() -> Auth.types.BaseAuthContext | None:
     return AuthContext.get()
 
 
+def get_user_id(user: BaseUser | None) -> str | None:
+    if user is None:
+        return None
+    try:
+        return user.identity
+    except NotImplementedError:
+        try:
+            return user.display_name
+        except NotImplementedError:
+            pass
+
+
+def merge_auth(
+    config: RunnableConfig,
+    ctx: Auth.types.BaseAuthContext | None = None,
+) -> RunnableConfig:
+    """Inject auth context into config's configurable dict.
+
+    If ctx is not provided, attempts to get it from the current context.
+    """
+    if ctx is None:
+        ctx = get_auth_ctx()
+    if ctx is None:
+        return config
+
+    configurable = config.setdefault("configurable", {})
+    return config | {
+        "configurable": configurable
+        | {
+            "langgraph_auth_user": cast("BaseUser | None", ctx.user),
+            "langgraph_auth_user_id": get_user_id(cast("BaseUser | None", ctx.user)),
+            "langgraph_auth_permissions": ctx.permissions,
+        }
+    }
+
+
 class AsyncCursorProto(Protocol):
     async def fetchone(self) -> Row: ...
 
@@ -139,9 +176,17 @@ class SchemaGenerator(BaseSchemaGenerator):
         for endpoint in endpoints_info:
             try:
                 parsed = self.parse_docstring(endpoint.func)
-            except AssertionError:
-                logger.warning("Could not parse docstrings for route %s", endpoint.path)
-                parsed = {}
+            except Exception as exc:
+                docstring = getattr(endpoint.func, "__doc__", None) or ""
+                logger.warning(
+                    "Unable to parse docstring from OpenAPI schema for route %s (%s): %s\n\nUsing as description",
+                    endpoint.path,
+                    endpoint.func.__qualname__,
+                    exc,
+                    exc_info=exc,
+                    docstring=docstring,
+                )
+                parsed = {"description": docstring}
 
             if endpoint.path not in schema["paths"]:
                 schema["paths"][endpoint.path] = {}
@@ -186,14 +231,14 @@ def validate_select_columns(
 
 
 __all__ = [
+    "AsyncConnectionProto",
     "AsyncCursorProto",
     "AsyncPipelineProto",
-    "AsyncConnectionProto",
-    "fetchone",
-    "validate_uuid",
-    "next_cron_date",
     "SchemaGenerator",
+    "fetchone",
     "get_pagination_headers",
+    "next_cron_date",
     "uuid7",
     "validate_select_columns",
+    "validate_uuid",
 ]

langgraph_api/utils/config.py

@@ -4,7 +4,6 @@ import asyncio
 import functools
 import typing
 from collections import ChainMap
-from concurrent.futures import Executor
 from contextvars import copy_context
 from os import getenv
 from typing import Any, ParamSpec, TypeVar
@@ -13,6 +12,8 @@ from langgraph.constants import CONF
 from typing_extensions import TypedDict
 
 if typing.TYPE_CHECKING:
+    from concurrent.futures import Executor
+
     from langchain_core.runnables import RunnableConfig
 
 try:

langgraph_api/utils/future.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
-import contextvars
 import inspect
 import sys
 import types
-from collections.abc import Awaitable, Coroutine, Generator
-from typing import TypeVar, cast
+from typing import TYPE_CHECKING, TypeVar, cast
+
+if TYPE_CHECKING:
+    import contextvars
+    from collections.abc import Awaitable, Coroutine, Generator
 
 T = TypeVar("T")
 AnyFuture = asyncio.Future | concurrent.futures.Future
@@ -45,7 +47,8 @@ def _set_concurrent_future_state(
     source: AnyFuture,
 ) -> None:
     """Copy state from a future to a concurrent.futures.Future."""
-    assert source.done()
+    if not source.done():
+        raise ValueError("Future is not done")
     if source.cancelled():
         concurrent.cancel()
     if not concurrent.set_running_or_notify_cancel():
@@ -65,7 +68,8 @@ def _copy_future_state(source: AnyFuture, dest: asyncio.Future) -> None:
     """
     if dest.done():
         return
-    assert source.done()
+    if not source.done():
+        raise ValueError("Future is not done")
     if dest.cancelled():
         return
     if source.cancelled():
@@ -152,7 +156,7 @@ def _ensure_future(
     if not asyncio.iscoroutine(coro_or_future):
         if inspect.isawaitable(coro_or_future):
             coro_or_future = cast(
-                Coroutine[None, None, T], _wrap_awaitable(coro_or_future)
+                "Coroutine[None, None, T]", _wrap_awaitable(coro_or_future)
             )
             called_wrap_awaitable = True
         else:

langgraph_api/utils/uuids.py

@@ -1,32 +1,35 @@
-import os
-import time
-from uuid import UUID, SafeUUID
+"""UUID utility functions.
 
-_last_timestamp_v7 = None
-_last_counter_v7 = 0  # 42-bit counter
-_RFC_4122_VERSION_7_FLAGS = (7 << 76) | (0x8000 << 48)
+This module exports a uuid7 function to generate monotonic, time-ordered UUIDs
+for tracing and similar operations.
+"""
 
+from __future__ import annotations
 
-def _uuid7_get_counter_and_tail():
-    rand = int.from_bytes(os.urandom(10))
-    # 42-bit counter with MSB set to 0
-    counter = (rand >> 32) & 0x1FF_FFFF_FFFF
-    # 32-bit random data
-    tail = rand & 0xFFFF_FFFF
-    return counter, tail
+import typing
+from uuid import UUID
 
+from uuid_utils.compat import uuid7 as _uuid_utils_uuid7
 
-def _from_int(value: int) -> UUID:
-    uid = object.__new__(UUID)
-    object.__setattr__(uid, "int", value)
-    object.__setattr__(uid, "is_safe", SafeUUID.unknown)
-    return uid
+if typing.TYPE_CHECKING:
+    from uuid import UUID
 
+_NANOS_PER_SECOND: typing.Final = 1_000_000_000
 
-def uuid7():
-    """Generate a UUID from a Unix timestamp in milliseconds and random bits.
+
+def _to_timestamp_and_nanos(nanoseconds: int) -> tuple[int, int]:
+    """Split a nanosecond timestamp into seconds and remaining nanoseconds."""
+    seconds, nanos = divmod(nanoseconds, _NANOS_PER_SECOND)
+    return seconds, nanos
+
+
+def uuid7(nanoseconds: int | None = None) -> UUID:
+    """Generate a UUID from a Unix timestamp in nanoseconds and random bits.
 
     UUIDv7 objects feature monotonicity within a millisecond.
+
+    Args:
+        nanoseconds: Optional ns timestamp. If not provided, uses current time.
     """
     # --- 48 ---   -- 4 --   --- 12 ---   -- 2 --   --- 30 ---   - 32 -
     # unix_ts_ms | version | counter_hi | variant | counter_lo | random
@@ -41,47 +44,11 @@ def uuid7():
     # advanced and the counter is reset to a random 42-bit integer with MSB
     # set to 0.
 
-    global _last_timestamp_v7
-    global _last_counter_v7
-
-    nanoseconds = time.time_ns()
-    timestamp_ms = nanoseconds // 1_000_000
-
-    if _last_timestamp_v7 is None or timestamp_ms > _last_timestamp_v7:
-        counter, tail = _uuid7_get_counter_and_tail()
-    else:
-        if timestamp_ms < _last_timestamp_v7:
-            timestamp_ms = _last_timestamp_v7 + 1
-        # advance the 42-bit counter
-        counter = _last_counter_v7 + 1
-        if counter > 0x3FF_FFFF_FFFF:
-            # advance the 48-bit timestamp
-            timestamp_ms += 1
-            counter, tail = _uuid7_get_counter_and_tail()
-        else:
-            # 32-bit random data
-            tail = int.from_bytes(os.urandom(4))
-
-    unix_ts_ms = timestamp_ms & 0xFFFF_FFFF_FFFF
-    counter_msbs = counter >> 30
-    # keep 12 counter's MSBs and clear variant bits
-    counter_hi = counter_msbs & 0x0FFF
-    # keep 30 counter's LSBs and clear version bits
-    counter_lo = counter & 0x3FFF_FFFF
-    # ensure that the tail is always a 32-bit integer (by construction,
-    # it is already the case, but future interfaces may allow the user
-    # to specify the random tail)
-    tail &= 0xFFFF_FFFF
+    # For now, just delegate to the uuid_utils implementation
+    if nanoseconds is None:
+        return _uuid_utils_uuid7()
+    seconds, nanos = _to_timestamp_and_nanos(nanoseconds)
+    return _uuid_utils_uuid7(timestamp=seconds, nanos=nanos)
 
-    int_uuid_7 = unix_ts_ms << 80
-    int_uuid_7 |= counter_hi << 64
-    int_uuid_7 |= counter_lo << 32
-    int_uuid_7 |= tail
-    # by construction, the variant and version bits are already cleared
-    int_uuid_7 |= _RFC_4122_VERSION_7_FLAGS
-    res = _from_int(int_uuid_7)
 
-    # defer global update until all computations are done
-    _last_timestamp_v7 = timestamp_ms
-    _last_counter_v7 = counter
-    return res
+__all__ = ["uuid7"]

langgraph_api/validation.py

@@ -52,6 +52,9 @@ ThreadCreate = jsonschema_rs.validator_for(
 ThreadPatch = jsonschema_rs.validator_for(
     openapi["components"]["schemas"]["ThreadPatch"]
 )
+ThreadPruneRequest = jsonschema_rs.validator_for(
+    openapi["components"]["schemas"]["ThreadPruneRequest"]
+)
 ThreadStateUpdate = jsonschema_rs.validator_for(
     {
         **openapi["components"]["schemas"]["ThreadStateUpdate"],
@@ -124,6 +127,9 @@ RunCreateStateful = jsonschema_rs.validator_for(
 )
 RunsCancel = jsonschema_rs.validator_for(openapi["components"]["schemas"]["RunsCancel"])
 CronCreate = jsonschema_rs.validator_for(openapi["components"]["schemas"]["CronCreate"])
+ThreadCronCreate = jsonschema_rs.validator_for(
+    openapi["components"]["schemas"]["ThreadCronCreate"]
+)
 CronSearch = jsonschema_rs.validator_for(openapi["components"]["schemas"]["CronSearch"])
 CronCountRequest = jsonschema_rs.validator_for(
     openapi["components"]["schemas"]["CronCountRequest"]