langgraph-executor 0.0.1a0__py3-none-any.whl

langgraph_executor/executor.py

@@ -0,0 +1,341 @@
+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
+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.extract_graph import extract_graph
+from langgraph_executor.pb import executor_pb2, executor_pb2_grpc, types_pb2
+from langgraph_executor.stream_utils import ExecutorStreamHandler
+
+
+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.
+
+        Args:
+            graphs: Dictionary mapping graph names to compiled graphs
+
+        """
+        self.graphs = graphs
+        self.logger = logging.getLogger(__name__)
+
+    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()),
+        )
+
+    def GetGraph(self, request: Any, context: Any) -> executor_pb2.GetGraphResponse:  # type: ignore[name-defined]
+        """Get graph definition."""
+        try:
+            self.logger.debug("GetGraph called")
+
+            graph = self.graphs[request.graph_name]
+
+            # extract graph
+            graph_definition = extract_graph(graph)
+
+            return executor_pb2.GetGraphResponse(graph_definition=graph_definition)
+
+        except Exception as e:
+            self.logger.error(f"GetGraph Error: {e}", exc_info=True)
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    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)
+
+            # reconstruct specs
+            specs, _ = reconstruct_channels(
+                request.specs.channels,
+                graph,
+                scratchpad=None,  # type: ignore[invalid-arg-type]
+            )
+
+            # 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,
+                ),
+            )
+
+        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)
+                ]
+            )
+
+    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)
+
+
+@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
+
+    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())
+
+    BaseMessage.__init__ = patched_init  # type: ignore[method-assign]

langgraph_executor/extract_graph.py

@@ -0,0 +1,178 @@
+"""Shared module for extracting graph information from LangGraph graphs."""
+
+from collections.abc import Sequence
+from typing import Any
+
+from google.protobuf.json_format import MessageToJson
+from google.protobuf.struct_pb2 import Struct  # type: ignore[import-not-found]
+from langchain_core.runnables import RunnableConfig
+from langgraph._internal._constants import (  # CONFIG_KEY_PREVIOUS,
+    CONFIG_KEY_CHECKPOINT_ID,
+    CONFIG_KEY_CHECKPOINT_MAP,
+    CONFIG_KEY_CHECKPOINT_NS,
+    CONFIG_KEY_DURABILITY,
+    CONFIG_KEY_RESUMING,
+    CONFIG_KEY_TASK_ID,
+    CONFIG_KEY_THREAD_ID,
+    RESERVED,
+)
+from langgraph.cache.memory import InMemoryCache
+from langgraph.pregel import Pregel
+from langgraph.pregel._read import PregelNode
+from langgraph.utils.config import ensure_config
+
+from langgraph_executor.common import extract_channels
+from langgraph_executor.pb import graph_pb2, types_pb2
+
+DEFAULT_MAX_CONCURRENCY = 1
+
+
+def extract_cache_type(cache: Any) -> str:
+    """Extract cache type from a cache object."""
+    if cache is None:
+        return "unsupported"
+    if isinstance(cache, InMemoryCache):
+        return "inMemory"
+    return "unsupported"
+
+
+def extract_config(config: RunnableConfig) -> types_pb2.RunnableConfig:
+    ensured_config = ensure_config(config)
+    # metadata
+    metadata_proto = Struct()
+    metadata = {k: v for k, v in ensured_config["metadata"].items()}
+    metadata_proto.update(metadata)
+    # configurable
+    configurable_proto = Struct()
+
+    configurable = {}
+    for k, v in ensured_config["configurable"].items():
+        if k not in RESERVED:
+            configurable[k] = v
+
+    configurable_proto.update(configurable)
+    return 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)
+        ),
+        metadata=metadata_proto,
+        configurable=configurable_proto,
+        reserved_configurable=extract_reserved_configurable(
+            config.get("configurable", {})
+        ),
+    )
+
+
+def extract_reserved_configurable(
+    configurable: dict[str, Any],
+) -> types_pb2.ReservedConfigurable:
+    return types_pb2.ReservedConfigurable(
+        resuming=bool(configurable.get(CONFIG_KEY_RESUMING, False)),
+        task_id=str(configurable.get(CONFIG_KEY_TASK_ID, "")),
+        thread_id=str(configurable.get(CONFIG_KEY_THREAD_ID, "")),
+        checkpoint_map=dict(configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {})),
+        checkpoint_id=str(configurable.get(CONFIG_KEY_CHECKPOINT_ID, "")),
+        checkpoint_ns=str(configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")),
+        durability=configurable.get(CONFIG_KEY_DURABILITY, "async"),
+    )
+
+
+def extract_nodes(nodes: dict[str, PregelNode]) -> dict[str, graph_pb2.NodeDefinition]:
+    out = {}
+    for k, v in nodes.items():
+        out[k] = extract_node(k, v)
+    return out
+
+
+def extract_node(name: str, node: PregelNode) -> graph_pb2.NodeDefinition:
+    if isinstance(node.channels, str):
+        channels = [node.channels]
+    elif isinstance(node.channels, list):
+        channels = node.channels
+    elif isinstance(node.channels, dict):
+        channels = [k for k, _ in node.channels.items()]
+    else:
+        channels = []
+    # TODO cache policy
+    return graph_pb2.NodeDefinition(
+        metadata=Struct(fields=node.metadata or {}),
+        name=name,
+        triggers=node.triggers,
+        tags=node.tags or [],
+        channels=channels,
+    )
+
+
+def extract_trigger_to_nodes(
+    trigger_to_nodes: dict[str, Sequence[str]] | Any,  # Allow Mapping type from graph
+) -> dict[str, graph_pb2.TriggerMapping]:
+    trigger_map = {}
+    for trigger, nodes in trigger_to_nodes.items():
+        if isinstance(nodes, dict) and "nodes" in nodes:
+            trigger_map[trigger] = graph_pb2.TriggerMapping(nodes=nodes["nodes"])
+        elif isinstance(nodes, list):
+            trigger_map[trigger] = graph_pb2.TriggerMapping(nodes=nodes)
+        else:
+            trigger_map[trigger] = graph_pb2.TriggerMapping(nodes=[])
+    return trigger_map
+
+
+def extract_graph(graph: Pregel) -> graph_pb2.GraphDefinition:
+    """Extract graph information from a compiled LangGraph graph.
+
+    Returns a protobuf message that contains all relevant orchestration information about the graph
+    """
+    # Handle input_channels and output_channels oneof
+    graph_def = graph_pb2.GraphDefinition(
+        name=str(graph.name),
+        channels=extract_channels(graph.channels),
+        interrupt_before_nodes=list(graph.interrupt_before_nodes),
+        interrupt_after_nodes=list(graph.interrupt_after_nodes),
+        stream_mode=(
+            [graph.stream_mode]
+            if isinstance(graph.stream_mode, str)
+            else graph.stream_mode
+        ),
+        stream_eager=bool(graph.stream_eager),
+        stream_channels=(
+            [graph.stream_channels]
+            if isinstance(graph.stream_channels, str)
+            else list(graph.stream_channels)
+            if graph.stream_channels
+            else []
+        ),
+        step_timeout=float(graph.step_timeout) if graph.step_timeout else 0.0,
+        debug=bool(graph.debug),
+        # TODO retry policy
+        cache=graph_pb2.Cache(
+            cache_type=extract_cache_type(getattr(graph, "cache", None)),
+        ),
+        config=extract_config(graph.config) if graph.config else None,
+        nodes=extract_nodes(graph.nodes),
+        trigger_to_nodes=extract_trigger_to_nodes(graph.trigger_to_nodes),
+    )
+
+    # Set input_channels and input_channels_is_list flag based on type
+    if isinstance(graph.input_channels, str):
+        graph_def.input_channels.extend([graph.input_channels])
+        graph_def.input_channels_is_list = False
+    else:
+        graph_def.input_channels.extend(graph.input_channels)
+        graph_def.input_channels_is_list = True
+
+    # Set output_channels and output_channels_is_list flag based on type
+    if isinstance(graph.output_channels, str):
+        graph_def.output_channels.extend([graph.output_channels])
+        graph_def.output_channels_is_list = False
+    else:
+        graph_def.output_channels.extend(graph.output_channels)
+        graph_def.output_channels_is_list = True
+
+    return graph_def
+
+
+def convert_to_json(proto):
+    return MessageToJson(proto)

langgraph_executor/info_logger.py

@@ -0,0 +1,111 @@
+import json
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class ExecutorInfo:
+    id: str
+    pid: int
+    port: int
+    status: str
+    start_time: float
+    end_time: float | None = None
+    error_message: str | None = None
+
+
+class ExecutorInfoLogger:
+    def __init__(self, log_dir: Path):
+        self.logs_dir = log_dir / "executor"
+        self.info_file_name = "info.json"
+
+    def create_executor_logs_dir(self, executor_id: str) -> Path:
+        executor_dir = self.logs_dir / executor_id
+        executor_dir.mkdir(parents=True, exist_ok=True)
+
+        return executor_dir
+
+    def write_executor_info(self, executor_info: ExecutorInfo):
+        executor_dir = self.create_executor_logs_dir(executor_info.id)
+        info_file = executor_dir / self.info_file_name
+
+        data = {
+            "id": executor_info.id,
+            "pid": executor_info.pid,
+            "port": executor_info.port,
+            "status": executor_info.status,
+            "start_time": executor_info.start_time,
+            "end_time": executor_info.end_time,
+            "error_message": executor_info.error_message,
+        }
+
+        with info_file.open("w") as f:
+            json.dump(data, f, indent=2)
+            f.flush()
+
+    def read_executor_info(self, executor_id: str) -> ExecutorInfo | None:
+        try:
+            executor_dir = self.logs_dir / executor_id
+            info_file = executor_dir / self.info_file_name
+
+            if not info_file.exists():
+                return None
+
+            with open(info_file) as f:
+                data = json.load(f)
+
+            if not data:
+                return None
+
+            return ExecutorInfo(
+                id=data["id"],
+                pid=data["pid"],
+                port=data["port"],
+                status=data["status"],
+                start_time=data["start_time"],
+                end_time=data.get("end_time"),
+                error_message=data.get("error_message"),
+            )
+
+        except Exception as e:
+            print(f"Failed to read executor info: {e}")
+            return None
+
+    def update_executor_info(
+        self,
+        executor_id: str,
+        status: str,
+        error_message: str | None = None,
+        end_time: float | None = None,
+    ) -> None:
+        info_file = self.logs_dir / executor_id / self.info_file_name
+
+        if not info_file.exists():
+            return
+
+        # Read current info
+        with open(info_file) as f:
+            current_info = json.load(f)
+
+        if not current_info:
+            return
+
+        # Update fields
+        current_info["status"] = status
+        if error_message:
+            current_info["error_message"] = error_message
+        if end_time:
+            current_info["end_time"] = end_time
+
+        # Write
+        with open(info_file, "w") as f:
+            json.dump(current_info, f)
+
+    def cleanup_logs(self) -> None:
+        if not self.logs_dir.exists():
+            return
+
+        for dir in self.logs_dir.iterdir():
+            if dir.is_dir():
+                shutil.rmtree(dir)