langgraph-api 0.4.48__py3-none-any.whl → 0.5.6__py3-none-any.whl

langgraph_api/__init__.py

@@ -1 +1 @@
-__version__ = "0.4.48"
+__version__ = "0.5.6"

langgraph_api/api/assistants.py

@@ -1,3 +1,4 @@
+from functools import partial
 from typing import Any
 from uuid import uuid4
 
@@ -15,7 +16,7 @@ from starlette.routing import BaseRoute
 from langgraph_api import store as api_store
 from langgraph_api.feature_flags import FF_USE_CORE_API, USE_RUNTIME_CONTEXT_API
 from langgraph_api.graph import get_assistant_id, get_graph
-from langgraph_api.grpc_ops.ops import Assistants as GrpcAssistants
+from langgraph_api.grpc.ops import Assistants as GrpcAssistants
 from langgraph_api.js.base import BaseRemotePregel
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
 from langgraph_api.schema import ASSISTANT_FIELDS
@@ -37,7 +38,7 @@ from langgraph_api.validation import (
     ConfigValidator,
 )
 from langgraph_runtime.checkpoint import Checkpointer
-from langgraph_runtime.database import connect
+from langgraph_runtime.database import connect as base_connect
 from langgraph_runtime.ops import Assistants
 from langgraph_runtime.retry import retry_db
 
@@ -45,6 +46,8 @@ logger = structlog.stdlib.get_logger(__name__)
 
 CrudAssistants = GrpcAssistants if FF_USE_CORE_API else Assistants
 
+connect = partial(base_connect, supports_core_api=FF_USE_CORE_API)
+
 EXCLUDED_CONFIG_SCHEMA = (
     "__pregel_checkpointer",
     "__pregel_store",
@@ -255,7 +258,7 @@ async def get_assistant_graph(
     assistant_id = get_assistant_id(str(request.path_params["assistant_id"]))
     validate_uuid(assistant_id, "Invalid assistant ID: must be a UUID")
     async with connect() as conn:
-        assistant_ = await Assistants.get(conn, assistant_id)
+        assistant_ = await CrudAssistants.get(conn, assistant_id)
         assistant = await fetchone(assistant_)
     config = json_loads(assistant["config"])
     configurable = config.setdefault("configurable", {})
@@ -312,43 +315,44 @@ async def get_assistant_subgraphs(
     assistant_id = request.path_params["assistant_id"]
     validate_uuid(assistant_id, "Invalid assistant ID: must be a UUID")
     async with connect() as conn:
-        assistant_ = await Assistants.get(conn, assistant_id)
+        assistant_ = await CrudAssistants.get(conn, assistant_id)
         assistant = await fetchone(assistant_)
-        config = json_loads(assistant["config"])
-        configurable = config.setdefault("configurable", {})
-        configurable.update(get_configurable_headers(request.headers))
-        async with get_graph(
-            assistant["graph_id"],
-            config,
-            checkpointer=Checkpointer(),
-            store=(await api_store.get_store()),
-        ) as graph:
-            namespace = request.path_params.get("namespace")
-
-            if isinstance(graph, BaseRemotePregel):
-                return ApiResponse(
-                    await graph.fetch_subgraphs(
+
+    config = json_loads(assistant["config"])
+    configurable = config.setdefault("configurable", {})
+    configurable.update(get_configurable_headers(request.headers))
+    async with get_graph(
+        assistant["graph_id"],
+        config,
+        checkpointer=Checkpointer(),
+        store=(await api_store.get_store()),
+    ) as graph:
+        namespace = request.path_params.get("namespace")
+
+        if isinstance(graph, BaseRemotePregel):
+            return ApiResponse(
+                await graph.fetch_subgraphs(
+                    namespace=namespace,
+                    recurse=request.query_params.get("recurse", "False")
+                    in ("true", "True"),
+                )
+            )
+
+        try:
+            return ApiResponse(
+                {
+                    ns: _graph_schemas(subgraph)
+                    async for ns, subgraph in graph.aget_subgraphs(
                         namespace=namespace,
                         recurse=request.query_params.get("recurse", "False")
                         in ("true", "True"),
                     )
-                )
-
-            try:
-                return ApiResponse(
-                    {
-                        ns: _graph_schemas(subgraph)
-                        async for ns, subgraph in graph.aget_subgraphs(
-                            namespace=namespace,
-                            recurse=request.query_params.get("recurse", "False")
-                            in ("true", "True"),
-                        )
-                    }
-                )
-            except NotImplementedError:
-                raise HTTPException(
-                    422, detail="The graph does not support visualization"
-                ) from None
+                }
+            )
+        except NotImplementedError:
+            raise HTTPException(
+                422, detail="The graph does not support visualization"
+            ) from None
 
 
 @retry_db
@@ -359,40 +363,40 @@ async def get_assistant_schemas(
     assistant_id = request.path_params["assistant_id"]
     validate_uuid(assistant_id, "Invalid assistant ID: must be a UUID")
     async with connect() as conn:
-        assistant_ = await Assistants.get(conn, assistant_id)
-        # TODO Implementa  cache so we can de-dent and release this connection.
+        assistant_ = await CrudAssistants.get(conn, assistant_id)
         assistant = await fetchone(assistant_)
-        config = json_loads(assistant["config"])
-        configurable = config.setdefault("configurable", {})
-        configurable.update(get_configurable_headers(request.headers))
-        async with get_graph(
-            assistant["graph_id"],
-            config,
-            checkpointer=Checkpointer(),
-            store=(await api_store.get_store()),
-        ) as graph:
-            if isinstance(graph, BaseRemotePregel):
-                schemas = await graph.fetch_state_schema()
-                return ApiResponse(
-                    {
-                        "graph_id": assistant["graph_id"],
-                        "input_schema": schemas.get("input"),
-                        "output_schema": schemas.get("output"),
-                        "state_schema": schemas.get("state"),
-                        "config_schema": schemas.get("config"),
-                        "context_schema": schemas.get("context"),
-                    }
-                )
-
-            schemas = _graph_schemas(graph)
 
+    config = json_loads(assistant["config"])
+    configurable = config.setdefault("configurable", {})
+    configurable.update(get_configurable_headers(request.headers))
+    async with get_graph(
+        assistant["graph_id"],
+        config,
+        checkpointer=Checkpointer(),
+        store=(await api_store.get_store()),
+    ) as graph:
+        if isinstance(graph, BaseRemotePregel):
+            schemas = await graph.fetch_state_schema()
             return ApiResponse(
                 {
                     "graph_id": assistant["graph_id"],
-                    **schemas,
+                    "input_schema": schemas.get("input"),
+                    "output_schema": schemas.get("output"),
+                    "state_schema": schemas.get("state"),
+                    "config_schema": schemas.get("config"),
+                    "context_schema": schemas.get("context"),
                 }
             )
 
+        schemas = _graph_schemas(graph)
+
+        return ApiResponse(
+            {
+                "graph_id": assistant["graph_id"],
+                **schemas,
+            }
+        )
+
 
 @retry_db
 async def patch_assistant(

langgraph_api/api/meta.py

@@ -86,6 +86,12 @@ async def meta_metrics(request: ApiRequest):
                         "# HELP lg_api_num_running_runs The number of runs currently running.",
                         "# TYPE lg_api_num_running_runs gauge",
                         f'lg_api_num_running_runs{{project_id="{metadata.PROJECT_ID}", revision_id="{metadata.HOST_REVISION_ID}"}} {queue_stats["n_running"]}',
+                        "# HELP lg_api_pending_runs_wait_time_max The maximum time a run has been pending, in seconds.",
+                        "# TYPE lg_api_pending_runs_wait_time_max gauge",
+                        f'lg_api_pending_runs_wait_time_max{{project_id="{metadata.PROJECT_ID}", revision_id="{metadata.HOST_REVISION_ID}"}} {queue_stats.get("pending_runs_wait_time_max_secs") or 0}',
+                        "# HELP lg_api_pending_runs_wait_time_med The median pending wait time across runs, in seconds.",
+                        "# TYPE lg_api_pending_runs_wait_time_med gauge",
+                        f'lg_api_pending_runs_wait_time_med{{project_id="{metadata.PROJECT_ID}", revision_id="{metadata.HOST_REVISION_ID}"}} {queue_stats.get("pending_runs_wait_time_med_secs") or 0}',
                     ]
                 )
         except Exception as e:

langgraph_api/api/threads.py

@@ -6,7 +6,7 @@ from starlette.responses import Response
 from starlette.routing import BaseRoute
 
 from langgraph_api.feature_flags import FF_USE_CORE_API
-from langgraph_api.grpc_ops.ops import Threads as GrpcThreads
+from langgraph_api.grpc.ops import Threads as GrpcThreads
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
 from langgraph_api.schema import THREAD_FIELDS, ThreadStreamMode
 from langgraph_api.sse import EventSourceResponse

langgraph_api/auth/custom.py

@@ -355,34 +355,39 @@ def _solve_fastapi_dependencies(
     }
 
     async def decorator(scope: dict, request: Request):
-        async with AsyncExitStack() as stack:
-            all_solved = await asyncio.gather(
-                *(
-                    solve_dependencies(
-                        request=request,
-                        dependant=dependent,
-                        async_exit_stack=stack,
-                        embed_body_fields=False,
+        async with AsyncExitStack() as request_stack:
+            scope["fastapi_inner_astack"] = request_stack
+            async with AsyncExitStack() as stack:
+                scope["fastapi_function_astack"] = stack
+                all_solved = await asyncio.gather(
+                    *(
+                        solve_dependencies(
+                            request=request,
+                            dependant=dependent,
+                            async_exit_stack=stack,
+                            embed_body_fields=False,
+                        )
+                        for dependent in dependents.values()
                     )
-                    for dependent in dependents.values()
                 )
-            )
-            all_injected = await asyncio.gather(
-                *(
-                    _run_async(dependent.call, solved.values, is_async)
-                    for dependent, solved in zip(
-                        dependents.values(), all_solved, strict=False
+                all_injected = await asyncio.gather(
+                    *(
+                        _run_async(dependent.call, solved.values, is_async)
+                        for dependent, solved in zip(
+                            dependents.values(), all_solved, strict=False
+                        )
                     )
                 )
-            )
-            kwargs = {
-                name: value
-                for name, value in zip(dependents.keys(), all_injected, strict=False)
-            }
-            other_params = _extract_arguments_from_scope(
-                scope, _param_names, request=request
-            )
-            return await fn(**(kwargs | other_params))
+                kwargs = {
+                    name: value
+                    for name, value in zip(
+                        dependents.keys(), all_injected, strict=False
+                    )
+                }
+                other_params = _extract_arguments_from_scope(
+                    scope, _param_names, request=request
+                )
+                return await fn(**(kwargs | other_params))
 
     return decorator
 

langgraph_api/config.py

@@ -128,6 +128,45 @@ class StoreConfig(TypedDict, total=False):
     ttl: TTLConfig
 
 
+class SerdeConfig(TypedDict, total=False):
+    """Configuration for the built-in serde, which handles checkpointing of state.
+
+    If omitted, no serde is set up (the object store will still be present, however)."""
+
+    allowed_json_modules: list[list[str]] | Literal[True] | None
+    """Optional. List of allowed python modules to de-serialize custom objects from.
+    
+    If provided, only the specified modules will be allowed to be deserialized.
+    If omitted, no modules are allowed, and the object returned will simply be a json object OR
+    a deserialized langchain object.
+    
+    Example:
+    {...
+        "serde": {
+            "allowed_json_modules": [
+                ["my_agent", "my_file", "SomeType"],
+            ]
+        }
+    }
+
+    If you set this to True, any module will be allowed to be deserialized.
+
+    Example:
+    {...
+        "serde": {
+            "allowed_json_modules": true
+        }
+    }
+    
+    """
+    pickle_fallback: bool
+    """Optional. Whether to allow pickling as a fallback for deserialization.
+    
+    If True, pickling will be allowed as a fallback for deserialization.
+    If False, pickling will not be allowed as a fallback for deserialization.
+    Defaults to True if not configured."""
+
+
 class CheckpointerConfig(TypedDict, total=False):
     """Configuration for the built-in checkpointer, which handles checkpointing of state.
 
@@ -140,6 +179,8 @@ class CheckpointerConfig(TypedDict, total=False):
     If provided, the checkpointer will apply TTL settings according to the configuration.
     If omitted, no TTL behavior is configured.
     """
+    serde: SerdeConfig | None
+    """Optional. Defines the configuration for how checkpoints are serialized."""
 
 
 class SecurityConfig(TypedDict, total=False):
@@ -240,6 +281,9 @@ REDIS_URI = env("REDIS_URI", cast=str)
 REDIS_CLUSTER = env("REDIS_CLUSTER", cast=bool, default=False)
 REDIS_MAX_CONNECTIONS = env("REDIS_MAX_CONNECTIONS", cast=int, default=2000)
 REDIS_CONNECT_TIMEOUT = env("REDIS_CONNECT_TIMEOUT", cast=float, default=10.0)
+REDIS_HEALTH_CHECK_INTERVAL = env(
+    "REDIS_HEALTH_CHECK_INTERVAL", cast=float, default=10.0
+)
 REDIS_KEY_PREFIX = env("REDIS_KEY_PREFIX", cast=str, default="")
 RUN_STATS_CACHE_SECONDS = env("RUN_STATS_CACHE_SECONDS", cast=int, default=60)
 
@@ -250,6 +294,13 @@ ALLOW_PRIVATE_NETWORK = env("ALLOW_PRIVATE_NETWORK", cast=bool, default=False)
 See https://developer.chrome.com/blog/private-network-access-update-2024-03
 """
 
+# gRPC client pool size for persistence server.
+GRPC_CLIENT_POOL_SIZE = env("GRPC_CLIENT_POOL_SIZE", cast=int, default=5)
+
+# Minimum payload size to use the dedicated thread pool for JSON parsing.
+# (Otherwise, the payload is parsed directly in the event loop.)
+JSON_THREAD_POOL_MINIMUM_SIZE_BYTES = 100 * 1024  # 100 KB
+
 HTTP_CONFIG = env("LANGGRAPH_HTTP", cast=_parse_schema(HttpConfig), default=None)
 STORE_CONFIG = env("LANGGRAPH_STORE", cast=_parse_schema(StoreConfig), default=None)
 
@@ -339,6 +390,11 @@ def _parse_thread_ttl(value: str | None) -> ThreadTTLConfig | None:
 CHECKPOINTER_CONFIG = env(
     "LANGGRAPH_CHECKPOINTER", cast=_parse_schema(CheckpointerConfig), default=None
 )
+SERDE: SerdeConfig | None = (
+    CHECKPOINTER_CONFIG["serde"]
+    if CHECKPOINTER_CONFIG and "serde" in CHECKPOINTER_CONFIG
+    else None
+)
 THREAD_TTL: ThreadTTLConfig | None = env(
     "LANGGRAPH_THREAD_TTL", cast=_parse_thread_ttl, default=None
 )
@@ -349,7 +405,6 @@ N_JOBS_PER_WORKER = env("N_JOBS_PER_WORKER", cast=int, default=10)
 BG_JOB_TIMEOUT_SECS = env("BG_JOB_TIMEOUT_SECS", cast=float, default=3600)
 
 FF_CRONS_ENABLED = env("FF_CRONS_ENABLED", cast=bool, default=True)
-FF_RICH_THREADS = env("FF_RICH_THREADS", cast=bool, default=True)
 FF_LOG_DROPPED_EVENTS = env("FF_LOG_DROPPED_EVENTS", cast=bool, default=False)
 FF_LOG_QUERY_AND_PARAMS = env("FF_LOG_QUERY_AND_PARAMS", cast=bool, default=False)
 

langgraph_api/graph.py

@@ -51,7 +51,7 @@ async def register_graph(
     description: str | None = None,
 ) -> None:
     """Register a graph."""
-    from langgraph_api.grpc_ops.ops import Assistants as AssistantsGrpc
+    from langgraph_api.grpc.ops import Assistants as AssistantsGrpc
     from langgraph_runtime.database import connect
     from langgraph_runtime.ops import Assistants as AssistantsRuntime
 

langgraph_api/{grpc_ops → grpc}/client.py

@@ -1,5 +1,6 @@
 """gRPC client wrapper for LangGraph persistence services."""
 
+import asyncio
 import os
 
 import structlog
@@ -10,6 +11,10 @@ from .generated.core_api_pb2_grpc import AdminStub, AssistantsStub, ThreadsStub
 logger = structlog.stdlib.get_logger(__name__)
 
 
+# Shared global client pool
+_client_pool: "GrpcClientPool | None" = None
+
+
 class GrpcClient:
     """gRPC client for LangGraph persistence services."""
 
@@ -90,3 +95,89 @@ class GrpcClient:
                 "Client not connected. Use async context manager or call connect() first."
             )
         return self._admin_stub
+
+
+class GrpcClientPool:
+    """Pool of gRPC clients for load distribution."""
+
+    def __init__(self, pool_size: int = 5, server_address: str | None = None):
+        self.pool_size = pool_size
+        self.server_address = server_address
+        self.clients: list[GrpcClient] = []
+        self._current_index = 0
+        self._init_lock = asyncio.Lock()
+        self._initialized = False
+
+    async def _initialize(self):
+        """Initialize the pool of clients."""
+        async with self._init_lock:
+            if self._initialized:
+                return
+
+            await logger.ainfo(
+                "Initializing gRPC client pool",
+                pool_size=self.pool_size,
+                server_address=self.server_address,
+            )
+
+            for _ in range(self.pool_size):
+                client = GrpcClient(server_address=self.server_address)
+                await client.connect()
+                self.clients.append(client)
+
+            self._initialized = True
+            await logger.ainfo(
+                f"gRPC client pool initialized with {self.pool_size} clients"
+            )
+
+    async def get_client(self) -> GrpcClient:
+        """Get next client using round-robin selection.
+
+        Round-robin without strict locking - slight races are acceptable
+        and result in good enough distribution under high load.
+        """
+        if not self._initialized:
+            await self._initialize()
+
+        idx = self._current_index % self.pool_size
+        self._current_index = idx + 1
+        return self.clients[idx]
+
+    async def close(self):
+        """Close all clients in the pool."""
+        if self._initialized:
+            await logger.ainfo(f"Closing gRPC client pool ({self.pool_size} clients)")
+            for client in self.clients:
+                await client.close()
+            self.clients.clear()
+            self._initialized = False
+
+
+async def get_shared_client() -> GrpcClient:
+    """Get a gRPC client from the shared pool.
+
+    Uses a pool of channels for better performance under high concurrency.
+    Each channel is a separate TCP connection that can handle ~100-200
+    concurrent streams effectively.
+
+    Returns:
+        A GrpcClient instance from the pool
+    """
+    global _client_pool
+    if _client_pool is None:
+        from langgraph_api import config
+
+        _client_pool = GrpcClientPool(
+            pool_size=config.GRPC_CLIENT_POOL_SIZE,
+            server_address=os.getenv("GRPC_SERVER_ADDRESS"),
+        )
+
+    return await _client_pool.get_client()
+
+
+async def close_shared_client():
+    """Close the shared gRPC client pool."""
+    global _client_pool
+    if _client_pool is not None:
+        await _client_pool.close()
+        _client_pool = None