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

langgraph_executor/__init__.py

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

langgraph_executor/execute_task.py

@@ -43,21 +43,13 @@ from langgraph_executor.common import (
 from langgraph_executor.pb import types_pb2
 
 
-def get_init_request(request_iterator):
-    request = next(request_iterator)
-
-    if not hasattr(request, "init"):
-        raise ValueError("First message must be init")
-
-    return request.init
-
-
 def reconstruct_task(
     request,
     graph: Pregel,
     *,
     store: BaseStore | None = None,
     config: RunnableConfig | None = None,
+    custom_stream_writer=None,
 ) -> PregelExecutableTask:
     pb_task = request.task
 
@@ -91,7 +83,9 @@ def reconstruct_task(
             val = pb_to_val(pb_task.input["PUSH_INPUT"])
 
         writes = deque()
-        runtime = ensure_runtime(configurable, store, graph)
+        runtime = ensure_runtime(
+            configurable, store, graph, custom_stream_writer=custom_stream_writer
+        )
 
         # Generate cache key if cache policy exists
         cache_policy = getattr(proc, "cache_policy", None)
@@ -184,15 +178,31 @@ def create_scratchpad(
     return scratchpad
 
 
-def ensure_runtime(configurable, store, graph):
+def ensure_runtime(configurable, store, graph, custom_stream_writer=None):
     runtime = configurable.get(CONFIG_KEY_RUNTIME)
+
+    # Prepare runtime overrides
+    overrides = {"store": store}
+    if custom_stream_writer is not None:
+        overrides["stream_writer"] = custom_stream_writer
+
     if runtime is None:
-        return DEFAULT_RUNTIME.override(store=store)
+        return DEFAULT_RUNTIME.override(**overrides)
     if isinstance(runtime, Runtime):
-        return runtime.override(store=store)
+        return runtime.override(**overrides)
     if isinstance(runtime, dict):
         context = _coerce_context(graph, runtime.get("context"))
-        return Runtime(**(runtime | {"store": store, "context": context}))
+        return Runtime(
+            **(
+                runtime
+                | {"store": store, "context": context}
+                | (
+                    {"stream_writer": custom_stream_writer}
+                    if custom_stream_writer
+                    else {}
+                )
+            )
+        )
     raise ValueError("Invalid runtime")
 
 

langgraph_executor/executor.py

@@ -1,341 +1,162 @@
+import contextlib
+import functools
 import logging
-import uuid
-from collections.abc import Iterator, Sequence
-from functools import lru_cache
 from typing import Any
 
 import grpc
-from google.protobuf.struct_pb2 import Struct  # type: ignore[import-untyped]
-from langchain_core.messages import BaseMessage, BaseMessageChunk
-from langgraph.checkpoint.base import Checkpoint
-from langgraph.errors import GraphBubbleUp, GraphInterrupt
+import grpc.aio
+from langgraph._internal._constants import NS_SEP
 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 run_with_retry
-
-from langgraph_executor.common import (
-    checkpoint_to_proto,
-    exception_to_pb,
-    extract_channels,
-    get_graph,
-    pb_to_val,
-    reconstruct_channels,
-    reconstruct_checkpoint,
-    reconstruct_task_writes,
-    updates_to_proto,
-)
-from langgraph_executor.execute_task import (
-    extract_writes,
-    get_init_request,
-    reconstruct_task,
+
+from langgraph_executor.executor_base import LangGraphExecutorServicer
+from langgraph_executor.pb.executor_pb2_grpc import (
+    add_LangGraphExecutorServicer_to_server,
 )
-from langgraph_executor.extract_graph import extract_graph
-from langgraph_executor.pb import executor_pb2, executor_pb2_grpc, types_pb2
-from langgraph_executor.stream_utils import ExecutorStreamHandler
 
+# Internal helpers
+LOGGER = logging.getLogger(__name__)
+
+
+def create_server(graphs: dict[str, Pregel], address: str) -> grpc.aio.Server:
+    graphs, subgraph_map = _load_graphs(graphs)
+    server = grpc.aio.server(
+        # Be permissive: allow client pings without active RPCs and accept intervals
+        # as low as 50s. Our clients still default to ~5m, but this avoids penalizing
+        # other, more frequent clients.
+        options=[
+            ("grpc.keepalive_permit_without_calls", 1),
+            ("grpc.http2.min_recv_ping_interval_without_data_ms", 50000),  # 50s
+            ("grpc.http2.max_ping_strikes", 2),
+        ]
+    )
+    getter = functools.partial(get_graph, graphs=graphs)
+    add_LangGraphExecutorServicer_to_server(
+        LangGraphExecutorServicer(graphs, subgraph_map=subgraph_map, get_graph=getter),
+        server,
+    )
+    server.add_insecure_port(address)
+    return server
 
-class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
-    """gRPC servicer for LangGraph runtime execution operations."""
 
-    def __init__(self, graphs: dict[str, Pregel]):
-        """Initialize the servicer with compiled graphs.
+@contextlib.asynccontextmanager
+async def get_graph(graph_name: str, config: Any, *, graphs: dict[str, Pregel]):
+    yield graphs[graph_name]
 
-        Args:
-            graphs: Dictionary mapping graph names to compiled graphs
 
-        """
-        self.graphs = graphs
-        self.logger = logging.getLogger(__name__)
+def _load_graphs(graphs: dict[str, Pregel]) -> tuple[dict[str, Pregel], dict[str, str]]:
+    """Load graphs and their subgraphs recursively in hierarchical order.
 
-    def ListGraphs(self, request: Any, context: Any) -> executor_pb2.ListGraphsResponse:  # type: ignore[name-defined]
-        """List available graphs."""
-        return executor_pb2.ListGraphsResponse(
-            graph_names=list(self.graphs.keys()),
-        )
+    Args:
+        graphs: Dictionary of root graphs to load
+    """
+    # First, ensure all root graphs have unique names
+    _ensure_unique_root_names(graphs)
+    subgraph_map: dict[str, str] = {}
 
-    def GetGraph(self, request: Any, context: Any) -> executor_pb2.GetGraphResponse:  # type: ignore[name-defined]
-        """Get graph definition."""
-        try:
-            self.logger.debug("GetGraph called")
+    # Then, collect all subgraphs and mappings
+    all_subgraphs: dict[str, Pregel] = {}
+    subgraph_to_parent: dict[str, str] = {}
 
-            graph = self.graphs[request.graph_name]
+    for root_graph in graphs.values():
+        subgraphs, mappings = _collect_subgraphs(root_graph, root_graph.name)
+        all_subgraphs.update(subgraphs)
+        subgraph_to_parent.update(mappings)
 
-            # extract graph
-            graph_definition = extract_graph(graph)
+    subgraph_map.update(subgraph_to_parent)
 
-            return executor_pb2.GetGraphResponse(graph_definition=graph_definition)
+    # Now build self.graphs in hierarchical order (parents before children)
+    for root_name in sorted(graphs.keys()):
+        _load_graph_and_children(
+            root_name, graphs, {**graphs, **all_subgraphs}, subgraph_map
+        )
 
-        except Exception as e:
-            self.logger.error(f"GetGraph Error: {e}", exc_info=True)
-            context.abort(grpc.StatusCode.INTERNAL, str(e))
+    _log_supported_graphs(graphs, subgraph_map)
+    return graphs, subgraph_map
 
-    def ChannelsFromCheckpoint(
-        self, request: Any, context: Any
-    ) -> executor_pb2.ChannelsFromCheckpointResponse:  # type: ignore[name-defined]
-        try:
-            self.logger.debug("ChannelsFromCheckpoint called")
 
-            graph = get_graph(request.graph_name, self.graphs)
+def _ensure_unique_root_names(graphs: dict[str, Pregel]) -> None:
+    """Ensure all root graphs have unique names"""
+    seen_names = set()
 
-            # reconstruct specs
-            specs, _ = reconstruct_channels(
-                request.specs.channels,
-                graph,
-                scratchpad=None,  # type: ignore[invalid-arg-type]
+    for name in graphs:
+        if name in seen_names:
+            raise ValueError(
+                f"Root graph name conflict detected: {name}. Root graphs must have unique names"
             )
+        seen_names.add(name)
 
-            # initialize channels from specs and checkpoint channel values
-            checkpoint_dummy = Checkpoint(  # type: ignore[typeddict-item]
-                channel_values={
-                    k: pb_to_val(v)
-                    for k, v in request.checkpoint_channel_values.items()
-                },
-            )
-            channels, _ = channels_from_checkpoint(specs, checkpoint_dummy)
-
-            # channels to pb
-            channels = extract_channels(channels)
-
-            return executor_pb2.ChannelsFromCheckpointResponse(channels=channels)
-
-        except Exception as e:
-            self.logger.error(f"ChannelsFromCheckpoint Error: {e}", exc_info=True)
-            context.abort(grpc.StatusCode.INTERNAL, str(e))
-
-    def ExecuteTask(
-        self,
-        request_iterator: Iterator[executor_pb2.ExecuteTaskRequest],  # type: ignore[name-defined]
-        context: Any,
-    ) -> Iterator[executor_pb2.ExecuteTaskResponse]:  # type: ignore[name-defined]
-        self.logger.debug("ExecuteTask called")
-        _patch_specific_base_message()
-
-        # Right now, only handle task execution without interrupts, etc
-        try:
-            # Get request
-            request = get_init_request(request_iterator)
-
-            # Reconstruct PregelExecutableTask
-            graph = get_graph(request.graph_name, self.graphs)
-            task = reconstruct_task(request, graph)
-
-            # Check if streaming is requested (for messages mode)
-            stream_messages = "messages" in request.stream_modes
-
-            # Set up streaming callback if needed
-            stream_chunks = []
-            if stream_messages:
-
-                def stream_callback(message: BaseMessageChunk, metadata: dict):
-                    """Callback to capture stream chunks and queue them."""
-                    try:
-                        stream_chunks.append(
-                            executor_pb2.ExecuteTaskResponse(
-                                message_or_message_chunk=extract_output_message(message)
-                            )
-                        )
-                    except Exception as e:
-                        self.logger.warning(
-                            f"Failed to create stream chunk: {e}", exc_info=True
-                        )
-
-                # Create and inject callback handler
-                stream_handler = ExecutorStreamHandler(stream_callback, task.id)
-
-                # Add handler to task config callbacks
-                if "callbacks" not in task.config:
-                    task.config["callbacks"] = []
-                task.config["callbacks"].append(stream_handler)  # type: ignore[union-attr]
-
-            # Execute task, catching interrupts
-            # Check cache if task has cache key - send request to Go orchestrator
-            should_execute = True
-            if task.cache_key:
-                self.logger.debug(
-                    f"Task {task.id} has cache key, sending cache check request to Go",
-                )
-
-                # Send cache check request to Go runtime
-                cache_check_request = executor_pb2.CacheCheckRequest(
-                    cache_namespace=list(task.cache_key.ns),
-                    cache_key=task.cache_key.key,
-                    ttl=task.cache_key.ttl,
-                )
-
-                yield executor_pb2.ExecuteTaskResponse(
-                    cache_check_request=cache_check_request,
-                )
-
-                # Wait for Go's response via the bidirectional stream
-                try:
-                    cache_response_request = next(request_iterator)
-                    if hasattr(cache_response_request, "cache_check_response"):
-                        cache_response = cache_response_request.cache_check_response
-                        should_execute = not cache_response.cache_hit
-                        self.logger.debug(
-                            f"Received cache response for task {task.id}: cache_hit={cache_response.cache_hit}",
-                        )
-                    else:
-                        self.logger.warning(
-                            f"Expected cache_check_response for task {task.id}, got unexpected message type",
-                        )
-                        should_execute = (
-                            True  # Default to execution if unexpected response
-                        )
-                except StopIteration:
-                    self.logger.warning(
-                        f"No cache response received for task {task.id}, defaulting to execution",
-                    )
-                    should_execute = True  # Default to execution if no response
-
-            # TODO patch retry policy
-            # TODO configurable to deal with _call and the functional api
-
-            exception_pb = None
-            if not should_execute:
-                # Skip execution but still send response
-                pass
-            try:
-                run_with_retry(
-                    task,
-                    retry_policy=None,
-                )
-
-                # Yield any accumulated stream chunks
-                yield from stream_chunks
-
-            except Exception as e:
-                if isinstance(e, GraphBubbleUp | GraphInterrupt):
-                    self.logger.info(f"Interrupt in task {task.id}: {e}")
-                else:
-                    self.logger.exception(
-                        f"Exception running task {task.id}: {e}\nTask: {task}\n\n",
-                        exc_info=True,
-                    )
-                exception_pb = exception_to_pb(e)
-
-            # Send final messages via message_chunk if they exist
-            final_messages = extract_output_messages(task.writes)
-            if final_messages:
-                for message in final_messages:
-                    yield executor_pb2.ExecuteTaskResponse(
-                        message_or_message_chunk=message
-                    )
-
-            # Extract and yield channel writes
-            writes_pb = extract_writes(task.writes)
-            task_result_pb = (
-                executor_pb2.TaskResult(error=exception_pb, writes=writes_pb)
-                if exception_pb
-                else executor_pb2.TaskResult(writes=writes_pb)
-            )
 
-            yield executor_pb2.ExecuteTaskResponse(task_result=task_result_pb)
-
-            # Generate streaming chunks
-            # for chunk in output_writes(task, request):
-            #     yield executor_pb2.ExecuteTaskResponse(stream_chunk=chunk)
-
-        except Exception as e:
-            self.logger.exception(f"ExecuteTask error: {e}")
-            context.abort(grpc.StatusCode.INTERNAL, str(e))
-
-    def ApplyWrites(
-        self, request: Any, context: Any
-    ) -> executor_pb2.ApplyWritesResponse:  # type: ignore[name-defined]
-        # get graph
-        self.logger.debug("ApplyWrites called")
-        try:
-            # Reconstruct python objects from proto
-            graph = get_graph(request.graph_name, self.graphs)
-            channels, _ = reconstruct_channels(
-                request.channels.channels,
-                graph,
-                # TODO: figure this out
-                scratchpad=None,  # type: ignore[invalid-arg-type]
-            )
-            checkpoint = reconstruct_checkpoint(request.checkpoint)
-            tasks = reconstruct_task_writes(request.tasks)
-
-            # apply writes
-            updated_channel_names_set = apply_writes(
-                checkpoint,
-                channels,
-                tasks,
-                lambda *args: request.next_version,
-                graph.trigger_to_nodes,
-            )
-            updated_channel_names = list(updated_channel_names_set)
-
-            # Reconstruct protos
-            updated_channels = extract_channels(channels)
-            checkpoint_proto = checkpoint_to_proto(checkpoint)
-
-            # Respond with updates
-            return executor_pb2.ApplyWritesResponse(
-                updates=updates_to_proto(
-                    checkpoint_proto,
-                    updated_channel_names,
-                    updated_channels,
-                ),
-            )
+def _collect_subgraphs(
+    graph: Pregel, namespace: str
+) -> tuple[dict[str, Pregel], dict[str, str]]:
+    """Recursively collect all subgraphs from a root graph"""
+    subgraphs = {}
+    mappings = {}
 
-        except Exception as e:
-            self.logger.exception(f"ApplyWrites error: {e}")
-            context.abort(grpc.StatusCode.INTERNAL, str(e))
-
-
-def extract_output_messages(writes: Sequence[Any]) -> list[types_pb2.Message]:  # type: ignore[name-defined]
-    messages = []
-    for write in writes:
-        # Not sure this check is right
-        if isinstance(write[1], BaseMessage):
-            messages.append(extract_output_message(write[1]))
-        elif isinstance(write[1], Sequence):
-            messages.extend(
-                [
-                    extract_output_message(w)
-                    for w in write[1]
-                    if isinstance(w, BaseMessage)
-                ]
-            )
+    for idx, (node_name, subgraph) in enumerate(graph.get_subgraphs(recurse=False)):
+        # Generate subgraph name
+        subgraph.name = f"{namespace}{NS_SEP}{node_name}{NS_SEP}{idx}"
 
-    return messages
-
-
-def extract_output_message(write: Any) -> types_pb2.Message:  # type: ignore[name-defined]
-    message = Struct()
-    message.update(
-        {
-            "is_streaming_chunk": False,
-            "message": {
-                "id": getattr(write, "id", None) or uuid.uuid4().hex,
-                "type": getattr(write, "type", None),
-                "content": str(getattr(write, "content", "") or ""),
-                "additional_kwargs": getattr(write, "additional_kwargs", {}),
-                "usage_metadata": getattr(write, "usage_metadata", {}),
-                "tool_calls": getattr(write, "tool_calls", []),
-                "tool_call_id": getattr(write, "tool_call_id", ""),
-                "tool_call_chunks": getattr(write, "tool_call_chunks", []),
-                "response_metadata": getattr(write, "response_metadata", {}),
-            },
-            "metadata": {},
-        }
-    )
-    return types_pb2.Message(payload=message)
+        # Add this subgraph
+        subgraphs[subgraph.name] = subgraph
+        mappings[subgraph.name] = graph.name
+
+        # Recursively process this subgraph's children
+        nested_subgraphs, nested_mappings = _collect_subgraphs(subgraph, namespace)
+
+        subgraphs.update(nested_subgraphs)
+        mappings.update(nested_mappings)
+
+    return subgraphs, mappings
+
+
+def _load_graph_and_children(
+    graph_name: str,
+    graphs: dict[str, Pregel],
+    all_graphs: dict[str, Pregel],
+    subgraph_map: dict[str, str],
+) -> None:
+    """Recursively add a graph and its children to self.graphs in order"""
+
+    # Add this graph to self.graphs (maintaining insertion order)
+    graphs[graph_name] = all_graphs[graph_name]
+
+    # Get direct children of this graph
+    children = [
+        child_name
+        for child_name, parent_name in subgraph_map.items()
+        if parent_name == graph_name
+    ]
+
+    # Add children in sorted order (for deterministic output)
+    for child_name in sorted(children):
+        _load_graph_and_children(child_name, graphs, all_graphs, subgraph_map)
+
+
+def _log_supported_graphs(
+    graphs: dict[str, Pregel], subgraph_map: dict[str, str]
+) -> None:
+    """Log the complete graph hierarchy in a tree-like format."""
+    LOGGER.info("Loaded graphs:")
 
+    # Get root graphs
+    root_graphs = {name for name in graphs if name not in subgraph_map}
 
-@lru_cache(maxsize=1)
-def _patch_specific_base_message() -> None:
-    """Patch the specific BaseMessage class used in your system."""
-    from langchain_core.messages import BaseMessage
+    for root_name in sorted(root_graphs):
+        LOGGER.info(f"  {root_name}")
+        _log_graph_children(root_name, subgraph_map, indent=2)
 
-    original_init = BaseMessage.__init__
 
-    def patched_init(self, content: Any, **kwargs: Any) -> None:
-        original_init(self, content, **kwargs)
-        if self.id is None:
-            self.id = str(uuid.uuid4())
+def _log_graph_children(
+    parent_name: str, subgraph_map: dict[str, str], *, indent: int = 0
+) -> None:
+    """Recursively log children of a graph with proper indentation."""
+    children = [
+        child for child, parent in subgraph_map.items() if parent == parent_name
+    ]
 
-    BaseMessage.__init__ = patched_init  # type: ignore[method-assign]
+    for child in sorted(children):
+        prefix = "  " * indent + "└─ "
+        LOGGER.info(f"{prefix}{child}")
+        # Recursively log this child's children
+        _log_graph_children(child, subgraph_map, indent=indent + 1)