langgraph-executor 0.0.1a6__tar.gz → 0.0.1a8__tar.gz

{langgraph_executor-0.0.1a6 → langgraph_executor-0.0.1a8}/.gitignore

@@ -49,3 +49,4 @@ logs/
 #runtime binary 
 langgraph_runtime_integration/bin/*
 *.egg-info/
+.gocache

{langgraph_executor-0.0.1a6 → langgraph_executor-0.0.1a8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langgraph-executor
-Version: 0.0.1a6
+Version: 0.0.1a8
 Summary: LangGraph python RPC server executable by the langgraph-go orchestrator.
 Requires-Python: >=3.11
 Requires-Dist: grpcio>=1.73.1

langgraph_executor-0.0.1a8/langgraph_executor/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.1a8"

{langgraph_executor-0.0.1a6 → langgraph_executor-0.0.1a8}/langgraph_executor/client/patch.py

@@ -9,8 +9,8 @@ from langgraph.runtime import get_runtime
 from langgraph.types import Interrupt
 from pydantic import ValidationError
 
+from langgraph_executor import serde
 from langgraph_executor.client.utils import (
-    SERDE,
     config_to_pb,
     context_to_pb,
     create_runopts_pb,
@@ -58,7 +58,7 @@ async def _ainvoke_wrapper(
     # subgraph names coerced when initializing executor
     graph_name = pregel_self.name
 
-    logger.info(f"SUBGRAPH INVOKE ENCOUNTERED: {graph_name}")
+    logger.info(f"SUBGRAPH AINVOKE ENCOUNTERED: {graph_name}")
 
     # TODO: Hacky way of retrieving runtime from runnable context
     if not context:
@@ -111,7 +111,7 @@ async def _ainvoke_wrapper(
                     for interrupt in graph_interrupt.interrupts:
                         interrupts.append(
                             Interrupt(
-                                value=SERDE.loads_typed(
+                                value=serde.get_serializer().loads_typed(
                                     (
                                         interrupt.value.base_value.method,
                                         interrupt.value.base_value.value,
@@ -239,7 +239,7 @@ def _invoke_wrapper(
                     for interrupt in graph_interrupt.interrupts:
                         interrupts.append(
                             Interrupt(
-                                value=SERDE.loads_typed(
+                                value=serde.get_serializer().loads_typed(
                                     (
                                         interrupt.value.base_value.method,
                                         interrupt.value.base_value.value,

{langgraph_executor-0.0.1a6 → langgraph_executor-0.0.1a8}/langgraph_executor/client/utils.py

@@ -14,19 +14,18 @@ from langgraph._internal._constants import (
     CONFIG_KEY_CHECKPOINT_MAP,
     CONFIG_KEY_CHECKPOINT_NS,
     CONFIG_KEY_DURABILITY,
+    CONFIG_KEY_RESUME_MAP,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
     CONFIG_KEY_THREAD_ID,
 )
-from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 from langgraph.pregel.debug import CheckpointMetadata
 from langgraph.types import StateSnapshot
 
+from langgraph_executor import serde
 from langgraph_executor.common import reconstruct_config, val_to_pb
 from langgraph_executor.pb import runtime_pb2, types_pb2
 
-SERDE = JsonPlusSerializer()
-
 
 def input_to_pb(input):
     return val_to_pb(None, input)
@@ -51,10 +50,11 @@ def maybe_update_reserved_configurable(
         reserved_configurable.checkpoint_id = str(value)
     elif key == CONFIG_KEY_CHECKPOINT_NS:
         reserved_configurable.checkpoint_ns = str(value)
-    # elif key == CONFIG_KEY_PREVIOUS:
-    #     serde = JsonPlusSerializer()
-    #     meth, ser = serde.dumps_typed(value)
-    #     reserved_configurable.previous = types_pb2.SerializedValue(method=meth, value=bytes(ser))
+    elif key == CONFIG_KEY_RESUME_MAP and value is not None:
+        resume_map = cast(dict[str, Any], value)
+        for k, v in resume_map.items():
+            pb_value = val_to_pb(None, v)
+            reserved_configurable.resume_map[k].CopyFrom(pb_value)
     elif key == CONFIG_KEY_DURABILITY:
         reserved_configurable.durability = str(value)
     else:
@@ -221,7 +221,7 @@ def _deser_vals(current_chunk):
     if not isinstance(current_chunk, dict):
         return current_chunk
     if set(current_chunk.keys()) == VAL_KEYS:
-        return SERDE.loads_typed(
+        return serde.get_serializer().loads_typed(
             (current_chunk["method"], base64.b64decode(current_chunk["value"]))
         )
     for k, v in current_chunk.items():

{langgraph_executor-0.0.1a6 → langgraph_executor-0.0.1a8}/langgraph_executor/common.py

@@ -1,8 +1,9 @@
 import traceback
+import uuid
 from collections.abc import Mapping, Sequence
 from collections.abc import Sequence as SequenceType
 from contextvars import ContextVar
-from typing import Any, cast
+from typing import Any, Literal, cast
 
 from google.protobuf.json_format import MessageToDict
 from langchain_core.runnables import RunnableConfig
@@ -26,7 +27,6 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
     PendingWrite,
 )
-from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 from langgraph.errors import GraphBubbleUp, GraphInterrupt
 from langgraph.managed.base import ManagedValue, ManagedValueMapping
 from langgraph.pregel import Pregel
@@ -34,6 +34,7 @@ from langgraph.pregel._algo import PregelTaskWrites
 from langgraph.pregel._read import PregelNode
 from langgraph.types import Command, Interrupt, Send
 
+from langgraph_executor import serde
 from langgraph_executor.pb import types_pb2
 
 var_child_runnable_config: ContextVar[RunnableConfig | None] = ContextVar(
@@ -44,8 +45,6 @@ var_child_runnable_config: ContextVar[RunnableConfig | None] = ContextVar(
 def map_reserved_configurable(
     reserved_configurable: types_pb2.ReservedConfigurable,
 ) -> dict[str, Any]:
-    # serde = JsonPlusSerializer()
-
     return {
         CONFIG_KEY_RESUMING: reserved_configurable.resuming,
         CONFIG_KEY_TASK_ID: reserved_configurable.task_id,
@@ -80,7 +79,7 @@ def reconstruct_config(pb_config: types_pb2.RunnableConfig) -> RunnableConfig:
         tags=list(pb_config.tags),
         metadata=MessageToDict(pb_config.metadata),
         run_name=pb_config.run_name,
-        run_id=pb_config.run_id,
+        run_id=cast(uuid.UUID, pb_config.run_id),
         max_concurrency=pb_config.max_concurrency,
         recursion_limit=pb_config.recursion_limit,
         configurable=configurable,
@@ -138,7 +137,6 @@ def reconstruct_checkpoint(request_checkpoint: types_pb2.Checkpoint) -> Checkpoi
 def reconstruct_task_writes(
     request_tasks: SequenceType[Any],
 ) -> SequenceType[PregelTaskWrites]:
-    # serde = JsonPlusSerializer()
     return [
         PregelTaskWrites(
             tuple(t.task_path),
@@ -187,11 +185,11 @@ def get_node(node_name: str, graph: Pregel, graph_name: str) -> PregelNode:
 
 
 def pb_to_val(value: types_pb2.Value) -> Any:
-    serde = JsonPlusSerializer()
-
     value_kind = value.WhichOneof("message")
     if value_kind == "base_value":
-        return serde.loads_typed((value.base_value.method, value.base_value.value))
+        return serde.get_serializer().loads_typed(
+            (value.base_value.method, value.base_value.value)
+        )
     if value_kind == "sends":
         sends = []
         for send in value.sends.sends:
@@ -313,9 +311,7 @@ def base_value_to_pb(value: Any) -> types_pb2.Value:
 
 
 def serialize_value(value: Any) -> types_pb2.SerializedValue:
-    serde = JsonPlusSerializer()
-
-    meth, ser_val = serde.dumps_typed(value)
+    meth, ser_val = serde.get_serializer().dumps_typed(value)
     return types_pb2.SerializedValue(method=meth, value=bytes(ser_val))
 
 
@@ -442,7 +438,7 @@ def reconstruct_checkpoint_metadata(
         return None
 
     return CheckpointMetadata(
-        source=metadata_pb.source,
+        source=cast(Literal["input", "loop", "update", "fork"], metadata_pb.source),
         step=metadata_pb.step,
-        parents=metadata_pb.parents,
+        parents=dict(metadata_pb.parents) or {},
     )

{langgraph_executor-0.0.1a6 → langgraph_executor-0.0.1a8}/langgraph_executor/executor_base.py

@@ -91,14 +91,16 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
         get_graph: GetGraph,
         runtime_channel: grpc.Channel | None = None,
         logger: Logger | None = None,
-        on_message: Callable[
-            [
-                BaseMessageChunk,
-                dict[str, Any],
-            ],
-            None,
-        ]
-        | None = None,
+        on_message: (
+            Callable[
+                [
+                    BaseMessageChunk,
+                    dict[str, Any],
+                ],
+                None,
+            ]
+            | None
+        ) = None,
         on_custom: Callable[[Any], None] | None = None,
         get_store: Callable[[], Awaitable[BaseStore]] | None = None,
     ):
@@ -143,7 +145,9 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
     ) -> executor_pb2.GetGraphResponse:  # type: ignore[name-defined]
         """Get graph definition."""
         try:
-            self.logger.debug("GetGraph called")
+            self.logger.debug(
+                "GetGraph called", extra={"graph_name": request.graph_name}
+            )
             graph_name: str = request.graph_name
             return await self._get_graph_definition(graph_name)
 
@@ -183,7 +187,14 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
         self, request: Any, context: grpc.aio.ServicerContext
     ) -> executor_pb2.ChannelsFromCheckpointResponse:  # type: ignore[name-defined]
         try:
-            self.logger.debug("ChannelsFromCheckpoint called")
+            self.logger.debug(
+                "ChannelsFromCheckpoint called",
+                extra={
+                    "graph_name": request.graph_name,
+                    "specs": request.specs,
+                    "checkpoint_channel_values": request.checkpoint_channel_values,
+                },
+            )
             async with self.get_graph(request.graph_name, RunnableConfig()) as graph:
                 # reconstruct specs
                 specs, _ = reconstruct_channels(
@@ -215,11 +226,16 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
         request_iterator: Iterator[executor_pb2.ExecuteTaskRequest],  # type: ignore[name-defined]
         context: grpc.aio.ServicerContext,
     ) -> AsyncIterator[executor_pb2.ExecuteTaskResponse]:  # type: ignore[name-defined]
-        self.logger.debug("ExecuteTask called")
-
-        # Right now, only handle task execution without interrupts, etc
         try:
             request = await _get_init_request(request_iterator)
+            self.logger.debug(
+                "ExecuteTask called",
+                extra={
+                    "graph_name": request.graph_name,
+                    "task": request.task,
+                    "stream_modes": request.stream_modes,
+                },
+            )
             config = reconstruct_config(request.task.config)
             store = await self.get_store() if self.get_store is not None else None
             async with self.get_graph(request.graph_name, config) as graph:
@@ -338,7 +354,15 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
         self, request: Any, context: grpc.aio.ServicerContext
     ) -> executor_pb2.ApplyWritesResponse:  # type: ignore[name-defined]
         # get graph
-        self.logger.debug("ApplyWrites called")
+        self.logger.debug(
+            "ApplyWrites called",
+            extra={
+                "graph_name": request.graph_name,
+                "tasks": request.tasks,
+                "channels": request.channels,
+                "checkpoint": request.checkpoint,
+            },
+        )
         try:
             async with self.get_graph(request.graph_name, RunnableConfig()) as graph:
                 channels, _ = reconstruct_channels(
@@ -382,7 +406,15 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
         request: executor_pb2.StateUpdateRequest,
         context: grpc.aio.ServicerContext,
     ) -> executor_pb2.TaskResult | None:
-        self.logger.debug("StateUpdate called")
+        self.logger.debug(
+            "StateUpdate called",
+            extra={
+                "graph_name": request.graph_name,
+                "node_name": request.node_name,
+                "task_id": request.task_id,
+                "values": request.values,
+            },
+        )
 
         try:
             async with self.get_graph(request.graph_name, RunnableConfig()) as graph:

{langgraph_executor-0.0.1a6 → langgraph_executor-0.0.1a8}/langgraph_executor/extract_graph.py

@@ -1,7 +1,7 @@
 """Shared module for extracting graph information from LangGraph graphs."""
 
 from collections.abc import Sequence
-from typing import Any
+from typing import Any, cast
 
 from google.protobuf.json_format import MessageToJson
 from google.protobuf.struct_pb2 import Struct  # type: ignore[import-not-found]
@@ -55,8 +55,8 @@ def extract_config(config: RunnableConfig) -> types_pb2.RunnableConfig:
         tags=[t for t in ensured_config["tags"]],
         recursion_limit=ensured_config["recursion_limit"],
         run_name=ensured_config.get("run_name", ""),
-        max_concurrency=(
-            ensured_config.get("max_concurrency", DEFAULT_MAX_CONCURRENCY)
+        max_concurrency=cast(
+            int, ensured_config.get("max_concurrency", DEFAULT_MAX_CONCURRENCY)
         ),
         metadata=metadata_proto,
         configurable=configurable_proto,

langgraph_executor-0.0.1a8/langgraph_executor/serde.py

@@ -0,0 +1,13 @@
+from langgraph.checkpoint.serde.base import SerializerProtocol
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+SERIALIZER: SerializerProtocol = JsonPlusSerializer()
+
+
+def set_serializer(serializer: SerializerProtocol) -> None:
+    global SERIALIZER
+    SERIALIZER = serializer
+
+
+def get_serializer() -> SerializerProtocol:
+    return SERIALIZER

langgraph_executor-0.0.1a6/langgraph_executor/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.0.1a6"