langgraph-executor 0.0.1a2__py3-none-any.whl → 0.0.1a4__py3-none-any.whl

langgraph_executor/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.1a2"
+__version__ = "0.0.1a4"

langgraph_executor/executor_base.py

@@ -3,7 +3,14 @@ import contextlib
 import functools
 import logging
 import uuid
-from collections.abc import AsyncIterator, Callable, Collection, Iterator, Sequence
+from collections.abc import (
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Collection,
+    Iterator,
+    Sequence,
+)
 from typing import Any, Protocol, cast
 
 import grpc
@@ -17,6 +24,7 @@ from langgraph.pregel import Pregel
 from langgraph.pregel._algo import apply_writes
 from langgraph.pregel._checkpoint import channels_from_checkpoint
 from langgraph.pregel._retry import arun_with_retry
+from langgraph.store.base import BaseStore
 from langgraph.types import PregelExecutableTask
 
 from langgraph_executor.common import (
@@ -65,6 +73,16 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
         subgraph_map: dict[str, str],
         get_graph: GetGraph,
         logger: Logger | 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,
     ):
         """Initialize the servicer with compiled graphs.
 
@@ -82,6 +100,9 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
         self.get_graph = get_graph
         _patch_base_message_with_ids()
         self._graph_definition_cache: dict[str, executor_pb2.GetGraphResponse] = {}
+        self.on_message = on_message
+        self.on_custom = on_custom
+        self.get_store = get_store
 
     async def ListGraphs(
         self, request: Any, context: grpc.aio.ServicerContext
@@ -174,6 +195,7 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
         try:
             request = await _get_init_request(request_iterator)
             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:
                 stream_messages = "messages" in request.stream_modes
                 stream_custom = "custom" in request.stream_modes
@@ -181,13 +203,18 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
                 stream_queue = asyncio.Queue()
 
                 custom_stream_writer = (
-                    _create_custom_stream_writer(stream_queue, self.logger)
+                    _create_custom_stream_writer(
+                        stream_queue, self.logger, on_custom=self.on_custom
+                    )
                     if stream_custom
                     else None
                 )
 
                 task = reconstruct_task(
-                    request, graph, custom_stream_writer=custom_stream_writer
+                    request,
+                    graph,
+                    custom_stream_writer=custom_stream_writer,
+                    store=store,
                 )
                 if stream_messages:
                     # Create and inject callback handler
@@ -196,6 +223,7 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
                             stream_callback,
                             logger=self.logger,
                             stream_queue=stream_queue,
+                            on_message=self.on_message,
                         ),
                         task.id,
                     )
@@ -362,13 +390,16 @@ async def _run_task(
 
 def stream_callback(
     message: BaseMessageChunk,
-    metadata: dict,
+    metadata: dict[str, Any],
     *,
     logger: Logger,
     stream_queue: asyncio.Queue[executor_pb2.ExecuteTaskResponse],
+    on_message: Callable[[BaseMessageChunk, dict[str, Any]], None] | None = None,
 ):
     """Callback to capture stream chunks and queue them."""
     try:
+        if on_message is not None:
+            on_message(message, metadata)
         stream_queue.put_nowait(
             executor_pb2.ExecuteTaskResponse(
                 message_or_message_chunk=_extract_output_message(message)
@@ -378,12 +409,19 @@ def stream_callback(
         logger.warning(f"Failed to create stream chunk: {e}", exc_info=True)
 
 
-def _create_custom_stream_writer(stream_queue: asyncio.Queue[Any], logger: Logger):
+def _create_custom_stream_writer(
+    stream_queue: asyncio.Queue[Any],
+    logger: Logger,
+    *,
+    on_custom: Callable[[Any], None] | None = None,
+):
     """Create a proper stream_writer function for custom mode (like langgraph does)."""
 
     def stream_writer(content):
         """Custom stream writer that creates CustomStreamEvent messages."""
         try:
+            if on_custom is not None:
+                on_custom(content)
             # Create payload struct (like langgraph does)
             payload = Struct()
             if isinstance(content, str):

{langgraph_executor-0.0.1a2.dist-info → langgraph_executor-0.0.1a4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langgraph-executor
-Version: 0.0.1a2
+Version: 0.0.1a4
 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.1a2.dist-info → langgraph_executor-0.0.1a4.dist-info}/RECORD

@@ -1,9 +1,9 @@
-langgraph_executor/__init__.py,sha256=3EXaFnXv1z-8AVCGrDM9MpCrpg8I15YhcHCsHzX3iGI,24
+langgraph_executor/__init__.py,sha256=uM_3HhsbqJiR7MV8Zx7ur-gqgdixAWQv-HmjMR8c9vY,24
 langgraph_executor/common.py,sha256=w75Bqbrj5LOtiWoOOdIi45yVB05xYXDAnokwg7MgDQE,13688
 langgraph_executor/example.py,sha256=TcfxgC9VfpZFliWnuVdJMllCDa8ji7vrSCRWnmdsUA8,900
 langgraph_executor/execute_task.py,sha256=_m-WyVODhShI29qzkK0DXSPoTY4E_izbuaT7F_fVZVw,7262
 langgraph_executor/executor.py,sha256=oUg6q8WWjAIwuh2ik76IbCPMKfmegT5ePZsWFnvR4WE,5296
-langgraph_executor/executor_base.py,sha256=OpFhZNd-FehaTkSlnmD4bw9cAQ3FF8nWO9mim5YL0Cg,18509
+langgraph_executor/executor_base.py,sha256=CP2C59PWi90NA3949rvaB7GPPMwSMh4pfLqY9Vb51JE,19597
 langgraph_executor/extract_graph.py,sha256=f7NNeZigRQlTAYg_-mV8eBgQ9q8kuThCDwLjRO6wReQ,6365
 langgraph_executor/info_logger.py,sha256=kgoVnMxm_aloHal89NHfH_Qb9eixrqD0dHmgnrm_vL4,3110
 langgraph_executor/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -27,6 +27,6 @@ langgraph_executor/pb/types_pb2.py,sha256=rUboWBthlLM-iJOmZZta-HSekwTNFuXfWLaYHe
 langgraph_executor/pb/types_pb2.pyi,sha256=tBX4LzmfWZ-IYscIeAEBTDC02GWWKEAQVB2JPNHChDc,40957
 langgraph_executor/pb/types_pb2_grpc.py,sha256=EPv87wCc-6BNJ2xTNcb9d3ictDerK5cBt7qhd7EmJiQ,886
 langgraph_executor/pb/types_pb2_grpc.pyi,sha256=Dl8kkjhqb6F1Kt24mcFg7ppish4iKVfjRiiBxEjsMMA,413
-langgraph_executor-0.0.1a2.dist-info/METADATA,sha256=-xr_G-Z1Huy6YMjYSJHYZZkPD_tu7Q-ZwDgLluOmkTo,433
-langgraph_executor-0.0.1a2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langgraph_executor-0.0.1a2.dist-info/RECORD,,
+langgraph_executor-0.0.1a4.dist-info/METADATA,sha256=5Ng1JslMYube7c36AvUoXtoWVj5iBeUbv0g7qdTuddY,433
+langgraph_executor-0.0.1a4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langgraph_executor-0.0.1a4.dist-info/RECORD,,