langgraph-api 0.5.0__tar.gz → 0.5.3__tar.gz

{langgraph_api-0.5.0 → langgraph_api-0.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langgraph-api
-Version: 0.5.0
+Version: 0.5.3
 Author-email: Will Fu-Hinthorn <will@langchain.dev>, Josh Rogers <josh@langchain.dev>, Parker Rule <parker@langchain.dev>
 License: Elastic-2.0
 License-File: LICENSE
@@ -13,7 +13,7 @@ Requires-Dist: httpx>=0.25.0
 Requires-Dist: jsonschema-rs<0.30,>=0.20.0
 Requires-Dist: langchain-core>=0.3.64
 Requires-Dist: langgraph-checkpoint<4,>=3
-Requires-Dist: langgraph-runtime-inmem<0.16.0,>=0.15.0
+Requires-Dist: langgraph-runtime-inmem<0.17.0,>=0.16.0
 Requires-Dist: langgraph-sdk>=0.2.0
 Requires-Dist: langgraph<2,>=0.4.10
 Requires-Dist: langsmith>=0.3.45

langgraph_api-0.5.3/langgraph_api/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.5.3"

{langgraph_api-0.5.0 → langgraph_api-0.5.3}/langgraph_api/api/assistants.py

@@ -1,3 +1,4 @@
+from functools import partial
 from typing import Any
 from uuid import uuid4
 
@@ -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-0.5.0 → langgraph_api-0.5.3}/langgraph_api/config.py

@@ -291,6 +291,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)
 
@@ -395,7 +402,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-0.5.0 → langgraph_api-0.5.3}/langgraph_api/grpc_ops/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

langgraph_api-0.5.3/langgraph_api/grpc_ops/config_conversion.py

@@ -0,0 +1,225 @@
+"""Conversion utils for the RunnableConfig."""
+
+# THIS IS DUPLICATED
+# TODO: WFH - Deduplicate with the executor logic by moving into a separate package
+# Sequencing in the next PR.
+from typing import Any, cast
+
+import orjson
+from langchain_core.runnables.config import RunnableConfig
+
+from langgraph_api.grpc_ops.generated import engine_common_pb2
+
+CONFIG_KEY_SEND = "__pregel_send"
+CONFIG_KEY_READ = "__pregel_read"
+CONFIG_KEY_RESUMING = "__pregel_resuming"
+CONFIG_KEY_TASK_ID = "__pregel_task_id"
+CONFIG_KEY_THREAD_ID = "thread_id"
+CONFIG_KEY_CHECKPOINT_MAP = "checkpoint_map"
+CONFIG_KEY_CHECKPOINT_ID = "checkpoint_id"
+CONFIG_KEY_CHECKPOINT_NS = "checkpoint_ns"
+CONFIG_KEY_SCRATCHPAD = "__pregel_scratchpad"
+CONFIG_KEY_DURABILITY = "__pregel_durability"
+CONFIG_KEY_GRAPH_ID = "graph_id"
+
+
+def _durability_to_proto(
+    durability: str,
+) -> engine_common_pb2.Durability:
+    match durability:
+        case "async":
+            return engine_common_pb2.Durability.ASYNC
+        case "sync":
+            return engine_common_pb2.Durability.SYNC
+        case "exit":
+            return engine_common_pb2.Durability.EXIT
+        case _:
+            raise ValueError(f"invalid durability: {durability}")
+
+
+def _durability_from_proto(
+    durability: engine_common_pb2.Durability,
+) -> str:
+    match durability:
+        case engine_common_pb2.Durability.ASYNC:
+            return "async"
+        case engine_common_pb2.Durability.SYNC:
+            return "sync"
+        case engine_common_pb2.Durability.EXIT:
+            return "exit"
+        case _:
+            raise ValueError(f"invalid durability: {durability}")
+
+
+def config_to_proto(
+    config: RunnableConfig,
+) -> engine_common_pb2.EngineRunnableConfig | None:
+    # Prepare kwargs for construction
+    if not config:
+        return None
+    cp = {**config}
+    pb_config = engine_common_pb2.EngineRunnableConfig()
+    for k, v in (cp.pop("metadata", None) or {}).items():
+        if k == "run_attempt":
+            pb_config.run_attempt = v
+        elif k == "run_id":
+            pb_config.server_run_id = str(v)
+        else:
+            pb_config.metadata_json[k] = orjson.dumps(v)
+    if run_name := cp.pop("run_name", None):
+        pb_config.run_name = run_name
+
+    if run_id := cp.pop("run_id", None):
+        pb_config.run_id = str(run_id)
+
+    if (max_concurrency := cp.pop("max_concurrency", None)) and isinstance(
+        max_concurrency, int
+    ):
+        pb_config.max_concurrency = max_concurrency
+
+    if (recursion_limit := cp.pop("recursion_limit", None)) and isinstance(
+        recursion_limit, int
+    ):
+        pb_config.recursion_limit = recursion_limit
+
+    # Handle collections after construction
+    if (tags := cp.pop("tags", None)) and isinstance(tags, list):
+        pb_config.tags.extend(tags)
+
+    if (configurable := cp.pop("configurable", None)) and isinstance(
+        configurable, dict
+    ):
+        _inject_configurable_into_proto(configurable, pb_config)
+    if cp:
+        pb_config.extra_json.update({k: orjson.dumps(v) for k, v in cp.items()})
+
+    return pb_config
+
+
+RESTRICTED_RESERVED_CONFIGURABLE_KEYS = {
+    CONFIG_KEY_SEND,
+    CONFIG_KEY_READ,
+    CONFIG_KEY_SCRATCHPAD,
+}
+
+
+def _inject_configurable_into_proto(
+    configurable: dict[str, Any], proto: engine_common_pb2.EngineRunnableConfig
+) -> None:
+    extra = {}
+    for key, value in configurable.items():
+        if key == CONFIG_KEY_RESUMING:
+            proto.resuming = bool(value)
+        elif key == CONFIG_KEY_TASK_ID:
+            proto.task_id = str(value)
+        elif key == CONFIG_KEY_THREAD_ID:
+            proto.thread_id = str(value)
+        elif key == CONFIG_KEY_CHECKPOINT_MAP:
+            proto.checkpoint_map.update(cast(dict[str, str], value))
+        elif key == CONFIG_KEY_CHECKPOINT_ID:
+            proto.checkpoint_id = str(value)
+        elif key == CONFIG_KEY_CHECKPOINT_NS:
+            proto.checkpoint_ns = str(value)
+        elif key == CONFIG_KEY_DURABILITY and value:
+            proto.durability = _durability_to_proto(value)
+        elif key not in RESTRICTED_RESERVED_CONFIGURABLE_KEYS:
+            extra[key] = value
+    if extra:
+        proto.extra_configurable_json.update(
+            {k: orjson.dumps(v) for k, v in extra.items()}
+        )
+
+
+def context_to_json_bytes(context: dict[str, Any] | Any) -> bytes | None:
+    """Convert context to JSON bytes for proto serialization."""
+    if context is None:
+        return None
+
+    # Convert dataclass or other objects to dict if needed
+    if hasattr(context, "__dict__") and not hasattr(context, "items"):
+        # Convert dataclass to dict
+        context_dict = context.__dict__
+    elif hasattr(context, "items"):
+        # Already a dict-like object
+        context_dict = dict(context)
+    else:
+        # Try to convert to dict using vars()
+        context_dict = vars(context) if hasattr(context, "__dict__") else {}
+
+    return orjson.dumps(context_dict)
+
+
+def config_from_proto(
+    config_proto: engine_common_pb2.EngineRunnableConfig | None,
+) -> RunnableConfig:
+    if not config_proto:
+        return RunnableConfig(tags=[], metadata={}, configurable={})
+
+    configurable = _configurable_from_proto(config_proto)
+
+    metadata = {}
+    for k, v in config_proto.metadata_json.items():
+        metadata[k] = orjson.loads(v)
+    if config_proto.HasField("run_attempt"):
+        metadata["run_attempt"] = config_proto.run_attempt
+    if config_proto.HasField("server_run_id"):
+        metadata["run_id"] = config_proto.server_run_id
+
+    config = RunnableConfig()
+    if config_proto.extra_json:
+        for k, v in config_proto.extra_json.items():
+            config[k] = orjson.loads(v)  # type: ignore[invalid-key]
+    if config_proto.tags:
+        config["tags"] = list(config_proto.tags)
+    if metadata:
+        config["metadata"] = metadata
+    if configurable:
+        config["configurable"] = configurable
+    if config_proto.HasField("run_name"):
+        config["run_name"] = config_proto.run_name
+
+    if config_proto.HasField("max_concurrency"):
+        config["max_concurrency"] = config_proto.max_concurrency
+
+    if config_proto.HasField("recursion_limit"):
+        config["recursion_limit"] = config_proto.recursion_limit
+
+    return config
+
+
+def _configurable_from_proto(
+    config_proto: engine_common_pb2.EngineRunnableConfig,
+) -> dict[str, Any]:
+    configurable = {}
+
+    if config_proto.HasField("resuming"):
+        configurable[CONFIG_KEY_RESUMING] = config_proto.resuming
+
+    if config_proto.HasField("task_id"):
+        configurable[CONFIG_KEY_TASK_ID] = config_proto.task_id
+
+    if config_proto.HasField("thread_id"):
+        configurable[CONFIG_KEY_THREAD_ID] = config_proto.thread_id
+
+    if config_proto.HasField("checkpoint_id"):
+        configurable[CONFIG_KEY_CHECKPOINT_ID] = config_proto.checkpoint_id
+
+    if config_proto.HasField("checkpoint_ns"):
+        configurable[CONFIG_KEY_CHECKPOINT_NS] = config_proto.checkpoint_ns
+
+    if config_proto.HasField("durability"):
+        durability = _durability_from_proto(config_proto.durability)
+        if durability:
+            configurable[CONFIG_KEY_DURABILITY] = durability
+
+    if config_proto.HasField("graph_id"):
+        configurable[CONFIG_KEY_GRAPH_ID] = config_proto.graph_id
+
+    if len(config_proto.checkpoint_map) > 0:
+        configurable[CONFIG_KEY_CHECKPOINT_MAP] = dict(config_proto.checkpoint_map)
+
+    if len(config_proto.extra_configurable_json) > 0:
+        for k, v in config_proto.extra_configurable_json.items():
+            configurable[k] = orjson.loads(v)
+
+    return configurable