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

langgraph_api/grpc/client.py

@@ -0,0 +1,183 @@
+"""gRPC client wrapper for LangGraph persistence services."""
+
+import asyncio
+import os
+
+import structlog
+from grpc import aio  # type: ignore[import]
+
+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."""
+
+    def __init__(
+        self,
+        server_address: str | None = None,
+    ):
+        """Initialize the gRPC client.
+
+        Args:
+            server_address: The gRPC server address (default: localhost:50051)
+        """
+        self.server_address = server_address or os.getenv(
+            "GRPC_SERVER_ADDRESS", "localhost:50051"
+        )
+        self._channel: aio.Channel | None = None
+        self._assistants_stub: AssistantsStub | None = None
+        self._threads_stub: ThreadsStub | None = None
+        self._admin_stub: AdminStub | None = None
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.close()
+
+    async def connect(self):
+        """Connect to the gRPC server."""
+        if self._channel is not None:
+            return
+
+        self._channel = aio.insecure_channel(self.server_address)
+
+        self._assistants_stub = AssistantsStub(self._channel)
+        self._threads_stub = ThreadsStub(self._channel)
+        self._admin_stub = AdminStub(self._channel)
+
+        await logger.adebug(
+            "Connected to gRPC server", server_address=self.server_address
+        )
+
+    async def close(self):
+        """Close the gRPC connection."""
+        if self._channel is not None:
+            await self._channel.close()
+            self._channel = None
+            self._assistants_stub = None
+            self._threads_stub = None
+            self._admin_stub = None
+            await logger.adebug("Closed gRPC connection")
+
+    @property
+    def assistants(self) -> AssistantsStub:
+        """Get the assistants service stub."""
+        if self._assistants_stub is None:
+            raise RuntimeError(
+                "Client not connected. Use async context manager or call connect() first."
+            )
+        return self._assistants_stub
+
+    @property
+    def threads(self) -> ThreadsStub:
+        """Get the threads service stub."""
+        if self._threads_stub is None:
+            raise RuntimeError(
+                "Client not connected. Use async context manager or call connect() first."
+            )
+        return self._threads_stub
+
+    @property
+    def admin(self) -> AdminStub:
+        """Get the admin service stub."""
+        if self._admin_stub is None:
+            raise RuntimeError(
+                "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/grpc/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.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