langgraph-executor 0.0.1a0__py3-none-any.whl

langgraph_executor/__init__.py

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

langgraph_executor/common.py

@@ -0,0 +1,395 @@
+import traceback
+from collections.abc import Mapping, Sequence
+from collections.abc import Sequence as SequenceType
+from typing import Any, cast
+
+from google.protobuf.json_format import MessageToDict
+from langchain_core.runnables import RunnableConfig
+from langgraph._internal._constants import (
+    CONFIG_KEY_CHECKPOINT_ID,
+    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,
+    TASKS,
+)
+from langgraph._internal._scratchpad import PregelScratchpad
+from langgraph._internal._typing import MISSING
+from langgraph.channels.base import BaseChannel, EmptyChannelError
+from langgraph.checkpoint.base import Checkpoint, 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
+from langgraph.pregel._algo import PregelTaskWrites
+from langgraph.pregel._read import PregelNode
+from langgraph.types import Command, Interrupt, Send
+
+from langgraph_executor.pb import types_pb2
+
+
+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,
+        CONFIG_KEY_THREAD_ID: reserved_configurable.thread_id,
+        CONFIG_KEY_CHECKPOINT_MAP: reserved_configurable.checkpoint_map,
+        CONFIG_KEY_CHECKPOINT_ID: reserved_configurable.checkpoint_id,
+        CONFIG_KEY_CHECKPOINT_NS: reserved_configurable.checkpoint_ns,
+        CONFIG_KEY_RESUME_MAP: {
+            k: pb_to_val(v) for k, v in reserved_configurable.resume_map.items()
+        }
+        if reserved_configurable.resume_map
+        else None,
+        # CONFIG_KEY_PREVIOUS: serde.loads_typed( TODO
+        #     (
+        #         reserved_configurable.previous.method,
+        #         reserved_configurable.previous.value,
+        #     )
+        # )
+        # if reserved_configurable.previous
+        # and reserved_configurable.previous.method != "missing"
+        # else None,
+        CONFIG_KEY_DURABILITY: reserved_configurable.durability,
+    }
+
+
+def reconstruct_config(pb_config: types_pb2.RunnableConfig) -> RunnableConfig:
+    configurable = MessageToDict(pb_config.configurable)
+    for k, v in map_reserved_configurable(pb_config.reserved_configurable).items():
+        if v or k not in configurable:
+            configurable[k] = v
+    return RunnableConfig(
+        tags=list(pb_config.tags),
+        metadata=MessageToDict(pb_config.metadata),
+        run_name=pb_config.run_name,
+        run_id=pb_config.run_id,
+        max_concurrency=pb_config.max_concurrency,
+        recursion_limit=pb_config.recursion_limit,
+        configurable=configurable,
+    )
+
+
+def revive_channel(channel: BaseChannel, channel_pb: types_pb2.Channel) -> BaseChannel:
+    val_pb = channel_pb.checkpoint_result
+    val = pb_to_val(val_pb)
+
+    return channel.copy().from_checkpoint(val)
+
+
+def reconstruct_channels(
+    channels_pb: dict[str, types_pb2.Channel],
+    graph: Pregel,
+    scratchpad: PregelScratchpad,
+) -> tuple[dict[str, BaseChannel], ManagedValueMapping]:
+    channels = {}
+    managed = {}
+    for k, v in graph.channels.items():
+        if isinstance(v, BaseChannel):
+            assert k in channels_pb
+            channels[k] = revive_channel(v, channels_pb[k])
+        elif isinstance(v, ManagedValue) and scratchpad is not None:  # managed values
+            managed[k] = v.get(scratchpad)
+        else:
+            raise NotImplementedError(f"Unrecognized channel value: {type(v)}")
+
+    return channels, managed
+
+
+def reconstruct_checkpoint(request_checkpoint: types_pb2.Checkpoint) -> Checkpoint:
+    channel_versions = dict(request_checkpoint.channel_versions)
+    versions_seen = {
+        k: dict(v.channel_versions) for k, v in request_checkpoint.versions_seen.items()
+    }
+    # channel_values = unwrap_channel_values(request_checkpoint.channel_values)
+
+    return Checkpoint(
+        v=request_checkpoint.v,
+        id=request_checkpoint.id,
+        channel_versions=channel_versions,
+        # channel_values=channel_values,
+        versions_seen=versions_seen,
+        ts=request_checkpoint.ts,
+    )
+
+
+def reconstruct_task_writes(
+    request_tasks: SequenceType[Any],
+) -> SequenceType[PregelTaskWrites]:
+    # serde = JsonPlusSerializer()
+    return [
+        PregelTaskWrites(
+            tuple(t.task_path),
+            t.name,
+            [(w.channel, pb_to_val(w.value)) for w in t.writes],
+            t.triggers,
+        )
+        for t in request_tasks
+    ]
+
+
+def checkpoint_to_proto(checkpoint: Checkpoint) -> types_pb2.Checkpoint:
+    checkpoint_proto = types_pb2.Checkpoint()
+    checkpoint_proto.channel_versions.update(checkpoint["channel_versions"])
+    for node, versions_dict in checkpoint["versions_seen"].items():
+        checkpoint_proto.versions_seen[node].channel_versions.update(versions_dict)
+
+    return checkpoint_proto
+
+
+def updates_to_proto(
+    checkpoint_proto: types_pb2.Checkpoint,
+    updated_channel_names: Sequence[str],
+    channels: types_pb2.Channels,
+) -> types_pb2.Updates:
+    return types_pb2.Updates(
+        checkpoint=checkpoint_proto,
+        updated_channels=updated_channel_names,
+        channels=channels,
+    )
+
+
+def get_graph(
+    graph_name: str,
+    graphs: dict[str, Pregel],
+) -> Pregel:
+    if graph_name not in graphs:
+        raise ValueError(f"Graph {graph_name} not supported")
+    return graphs[graph_name]
+
+
+def get_node(node_name: str, graph: Pregel, graph_name: str) -> PregelNode:
+    if node_name not in graph.nodes:
+        raise ValueError(f"Node {node_name} not found in graph {graph_name}")
+    return graph.nodes[node_name]
+
+
+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))
+    if value_kind == "sends":
+        sends = []
+        for send in value.sends.sends:
+            node = send.node
+            arg = pb_to_val(send.arg)
+            sends.append(Send(node, arg))
+        return sends
+    if value_kind == "missing":
+        return MISSING
+    if value_kind == "command":
+        graph, update, resume, goto = None, None, None, ()
+        if value.command.graph is not None:
+            graph = value.command.graph
+        if value.command.update is not None:
+            if (
+                isinstance(value.command.update, dict)
+                and len(value.command.update) == 1
+                and "__root__" in value.command.update
+            ):
+                update = pb_to_val(value.command.update["__root__"])
+            else:
+                update = {k: pb_to_val(v) for k, v in value.command.update.items()}
+        if value.command.resume:
+            which = value.command.resume.WhichOneof("message")
+            if which == "value":
+                resume = pb_to_val(value.command.resume.value)
+            else:
+                resume_map = {
+                    k: pb_to_val(v)
+                    for k, v in value.command.resume.values.values.items()
+                }
+                resume = resume_map
+        if value.command.gotos:
+            gotos = []
+            for g in value.command.gotos:
+                which = g.WhichOneof("message")
+                if which == "node_name":
+                    gotos.append(g.node_name.name)
+                else:
+                    gotos.append(Send(g.send.node, pb_to_val(g.send.arg)))
+            if len(gotos) == 1:
+                gotos = gotos[0]
+            goto = gotos
+        return Command(graph=graph, update=update, resume=resume, goto=goto)
+    raise NotImplementedError(f"Unrecognized value kind: {value_kind}")
+
+
+def send_to_pb(send: Send) -> types_pb2.Send:
+    return types_pb2.Send(
+        node=send.node,
+        arg=val_to_pb(TASKS if isinstance(send.arg, Send) else None, send.arg),
+    )
+
+
+def sends_to_pb(sends: list[Send]) -> types_pb2.Value:
+    if not sends:
+        return missing_to_pb()
+    pb = []
+    for send in sends:
+        pb.append(send_to_pb(send))
+
+    return types_pb2.Value(sends=types_pb2.Sends(sends=pb))
+
+
+def command_to_pb(cmd: Command) -> types_pb2.Value:
+    cmd_pb = types_pb2.Command()
+    if cmd.graph:
+        if not cmd.graph == Command.PARENT:
+            raise ValueError("command graph must be null or parent")
+        cmd_pb.graph = cmd.graph
+    if cmd.update:
+        if isinstance(cmd.update, dict):
+            cmd_pb.update.update({k: val_to_pb(None, v) for k, v in cmd.update.items()})
+        else:
+            cmd_pb.update.update({"__root__": val_to_pb(None, cmd.update)})
+    if cmd.resume:
+        if isinstance(cmd.resume, dict):
+            cmd_pb.resume.CopyFrom(resume_map_to_pb(cmd.resume))
+        else:
+            resume_val = types_pb2.Resume(value=val_to_pb(None, cmd.resume))
+            cmd_pb.resume.CopyFrom(resume_val)
+    if cmd.goto:
+        gotos = []
+        goto = cmd.goto
+        if isinstance(goto, list):
+            for g in goto:
+                gotos.append(goto_to_pb(g))
+        else:
+            gotos.append(goto_to_pb(cast(Send | str, goto)))
+        cmd_pb.gotos.extend(gotos)
+
+    return types_pb2.Value(command=cmd_pb)
+
+
+def resume_map_to_pb(resume: dict[str, Any] | Any) -> types_pb2.Resume:
+    vals = {k: val_to_pb(None, v) for k, v in resume.items()}
+    return types_pb2.Resume(values=types_pb2.InterruptValues(values=vals))
+
+
+def goto_to_pb(goto: Send | str) -> types_pb2.Goto:
+    if isinstance(goto, Send):
+        return types_pb2.Goto(send=send_to_pb(goto))
+    if isinstance(goto, str):
+        return types_pb2.Goto(node_name=types_pb2.NodeName(name=goto))
+    raise ValueError("goto must be send or node name")
+
+
+def missing_to_pb() -> types_pb2.Value:
+    pb = types_pb2.Value()
+    pb.missing.SetInParent()
+    return pb
+
+
+def base_value_to_pb(value: Any) -> types_pb2.Value:
+    serialized_value = serialize_value(value)
+
+    return types_pb2.Value(base_value=serialized_value)
+
+
+def serialize_value(value: Any) -> types_pb2.SerializedValue:
+    serde = JsonPlusSerializer()
+
+    meth, ser_val = serde.dumps_typed(value)
+    return types_pb2.SerializedValue(method=meth, value=bytes(ser_val))
+
+
+def val_to_pb(channel_name: str | None, value: Any) -> types_pb2.Value:
+    if channel_name == TASKS and value != MISSING:
+        if not isinstance(value, list):
+            if not isinstance(value, Send):
+                raise ValueError(
+                    "Task must be a Send object objects."
+                    f" Got type={type(value)} value={value}",
+                )
+            value = [value]
+        else:
+            for v in value:
+                if not isinstance(v, Send):
+                    raise ValueError(
+                        "Task must be a list of Send objects."
+                        f" Got types={[type(v) for v in value]} values={value}",
+                    )
+        return sends_to_pb(value)
+    if value == MISSING:
+        return missing_to_pb()
+    if isinstance(value, Command):
+        return command_to_pb(value)
+    return base_value_to_pb(value)
+
+
+def extract_channel(name: str, channel: BaseChannel) -> types_pb2.Channel:
+    try:
+        get_result = channel.get()
+    except EmptyChannelError:
+        get_result = MISSING
+
+    return types_pb2.Channel(
+        get_result=val_to_pb(name, get_result),
+        is_available_result=channel.is_available(),
+        checkpoint_result=val_to_pb(name, channel.checkpoint()),
+    )
+
+
+def extract_channels(
+    channels: Mapping[str, BaseChannel | type[ManagedValue]],
+) -> types_pb2.Channels:
+    pb = {}
+    for name, channel in channels.items():
+        if isinstance(channel, BaseChannel):
+            pb[name] = extract_channel(name, channel)
+    return types_pb2.Channels(channels=pb)
+
+
+def exception_to_pb(exc: Exception) -> types_pb2.ExecutorError:
+    executor_error_pb = None
+    if isinstance(exc, GraphInterrupt):
+        if exc.args[0]:
+            interrupts = [interrupt_to_pb(interrupt) for interrupt in exc.args[0]]
+            graph_interrupt_pb = types_pb2.GraphInterrupt(
+                interrupts=interrupts,
+                interrupts_serialized=serialize_value(
+                    exc.args[0] if len(exc.args[0]) != 1 else exc.args[0][0],
+                ),  # brittle fix
+            )
+        else:
+            graph_interrupt_pb = types_pb2.GraphInterrupt()
+        executor_error_pb = types_pb2.ExecutorError(graph_interrupt=graph_interrupt_pb)
+    elif isinstance(exc, GraphBubbleUp):
+        bubbleup_pb = types_pb2.GraphBubbleUp()
+        executor_error_pb = types_pb2.ExecutorError(graph_bubble_up=bubbleup_pb)
+    else:
+        base_error_pb = types_pb2.BaseError(
+            error_type=str(type(exc)),
+            error_message=str(exc),
+            error_serialized=serialize_value(exc),
+        )
+        executor_error_pb = types_pb2.ExecutorError(base_error=base_error_pb)
+        executor_error_pb.traceback = traceback.format_exc()
+
+    return executor_error_pb
+
+
+def interrupt_to_pb(interrupt: Interrupt) -> types_pb2.Interrupt:
+    return types_pb2.Interrupt(
+        value=val_to_pb(None, interrupt.value),
+        id=interrupt.id,
+    )
+
+
+def pb_to_pending_writes(
+    pb: SequenceType[types_pb2.PendingWrite],
+) -> list[PendingWrite] | None:
+    if not pb:
+        return None
+
+    return [(pw.task_id, pw.channel, pb_to_val(pw.value)) for pw in pb]

langgraph_executor/example.py

@@ -0,0 +1,29 @@
+import grpc
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+from langgraph_executor.pb import runtime_pb2, runtime_pb2_grpc, types_pb2
+
+RUNTIME_SERVER_ADDRESS = "localhost:50051"
+
+if __name__ == "__main__":
+    channel = grpc.insecure_channel(RUNTIME_SERVER_ADDRESS)
+    stub = runtime_pb2_grpc.LangGraphRuntimeStub(channel)
+
+    serde = JsonPlusSerializer()
+    input_raw = {"messages": ["hi"], "count": 0}
+    method, ser = serde.dumps_typed(input_raw)
+    input = types_pb2.SerializedValue(method=method, value=bytes(ser))
+
+    request = runtime_pb2.InvokeRequest(
+        graph_name="example",
+        input=input,
+        config=types_pb2.RunnableConfig(
+            recursion_limit=25,
+            max_concurrency=1,
+            reserved_configurable=types_pb2.ReservedConfigurable(),
+        ),
+    )
+
+    response = stub.Invoke(request)
+
+    print("response: ", response)

langgraph_executor/execute_task.py

@@ -0,0 +1,239 @@
+from collections import deque
+from dataclasses import is_dataclass
+from functools import partial
+from typing import Any
+
+from langchain_core.runnables import RunnableConfig
+from langgraph._internal._config import patch_config
+from langgraph._internal._constants import (
+    CACHE_NS_WRITES,
+    CONF,
+    CONFIG_KEY_CHECKPOINT_NS,
+    CONFIG_KEY_READ,
+    CONFIG_KEY_RESUME_MAP,
+    CONFIG_KEY_RUNTIME,
+    CONFIG_KEY_SCRATCHPAD,
+    CONFIG_KEY_SEND,
+    PULL,
+    PUSH,
+)
+from langgraph._internal._scratchpad import PregelScratchpad
+from langgraph.pregel import Pregel
+from langgraph.pregel._algo import (
+    PregelTaskWrites,
+    _proc_input,
+    _scratchpad,
+    local_read,
+)
+from langgraph.pregel._call import identifier
+from langgraph.runtime import DEFAULT_RUNTIME, Runtime
+from langgraph.store.base import BaseStore
+from langgraph.types import CacheKey, PregelExecutableTask
+from pydantic import BaseModel
+from xxhash import xxh3_128_hexdigest
+
+from langgraph_executor.common import (
+    get_node,
+    pb_to_pending_writes,
+    pb_to_val,
+    reconstruct_channels,
+    reconstruct_config,
+    val_to_pb,
+)
+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,
+) -> PregelExecutableTask:
+    pb_task = request.task
+
+    try:
+        proc = get_node(pb_task.name, graph, pb_task.graph_name)
+        if config is None:
+            config = reconstruct_config(pb_task.config)
+        configurable = config.get(CONF, {})
+
+        scratchpad = create_scratchpad(
+            config,
+            pb_task,
+            request.step,
+            request.stop,
+        )
+        channels, managed = reconstruct_channels(
+            request.channels.channels,
+            graph,
+            scratchpad,
+        )
+        if pb_task.task_path[0] == PULL:
+            val = _proc_input(
+                proc,
+                managed,
+                channels,
+                for_execution=True,
+                scratchpad=scratchpad,
+                input_cache=None,
+            )
+        elif pb_task.task_path[0] == PUSH:
+            val = pb_to_val(pb_task.input["PUSH_INPUT"])
+
+        writes = deque()
+        runtime = ensure_runtime(configurable, store, graph)
+
+        # Generate cache key if cache policy exists
+        cache_policy = getattr(proc, "cache_policy", None)
+        cache_key = None
+        if cache_policy:
+            args_key = cache_policy.key_func(
+                *([val] if not isinstance(val, list | tuple) else val),
+            )
+            cache_key = CacheKey(
+                (CACHE_NS_WRITES, identifier(proc.node) or "__dynamic__"),
+                xxh3_128_hexdigest(
+                    args_key.encode() if isinstance(args_key, str) else args_key,
+                ),
+                cache_policy.ttl,
+            )
+
+        task = PregelExecutableTask(
+            name=pb_task.name,
+            input=val,
+            proc=proc.node,
+            writes=writes,
+            config=patch_config(
+                config,
+                configurable={
+                    CONFIG_KEY_SEND: writes.extend,
+                    CONFIG_KEY_READ: partial(
+                        local_read,
+                        scratchpad,
+                        channels,
+                        managed,
+                        PregelTaskWrites(
+                            tuple(pb_task.task_path)[:3],
+                            pb_task.name,
+                            writes,
+                            pb_task.triggers,
+                        ),
+                    ),
+                    CONFIG_KEY_RUNTIME: runtime,
+                    CONFIG_KEY_SCRATCHPAD: scratchpad,
+                },
+            ),
+            triggers=pb_task.triggers,
+            id=pb_task.id,
+            path=pb_task.task_path,
+            retry_policy=proc.retry_policy or [],  # TODO support
+            cache_key=cache_key,  # TODO support
+            writers=proc.flat_writers,
+            subgraphs=proc.subgraphs,
+        )
+
+    except Exception as e:
+        raise e
+
+    return task
+
+
+def extract_writes(writes) -> list[types_pb2.Write]:
+    w = []
+    for channel, val in writes:
+        val_pb = val_to_pb(channel, val)
+        channel_write = types_pb2.Write(channel=channel, value=val_pb)
+        w.append(channel_write)
+
+    return w
+
+
+def create_scratchpad(
+    config: RunnableConfig,
+    pb_task: types_pb2.Task,
+    step: int,
+    stop: int,
+) -> PregelScratchpad:
+    task_checkpoint_ns: str = config[CONF].get(CONFIG_KEY_CHECKPOINT_NS) or ""
+    pending_writes = (
+        pb_to_pending_writes(pb_task.pending_writes)
+        if (hasattr(pb_task, "pending_writes") and len(pb_task.pending_writes) > 0)
+        else []
+    )
+
+    scratchpad = _scratchpad(
+        config[CONF].get(CONFIG_KEY_SCRATCHPAD),
+        pending_writes or [],
+        pb_task.id,
+        xxh3_128_hexdigest(task_checkpoint_ns.encode()),
+        config[CONF].get(CONFIG_KEY_RESUME_MAP),
+        step,
+        stop,
+    )
+
+    return scratchpad
+
+
+def ensure_runtime(configurable, store, graph):
+    runtime = configurable.get(CONFIG_KEY_RUNTIME)
+    if runtime is None:
+        return DEFAULT_RUNTIME.override(store=store)
+    if isinstance(runtime, Runtime):
+        return runtime.override(store=store)
+    if isinstance(runtime, dict):
+        context = _coerce_context(graph, runtime.get("context"))
+        return Runtime(**(runtime | {"store": store, "context": context}))
+    raise ValueError("Invalid runtime")
+
+
+def _coerce_context(graph: Pregel, context: Any) -> Any:
+    if context is None:
+        return None
+
+    context_schema = graph.context_schema
+    if context_schema is None:
+        return context
+
+    schema_is_class = issubclass(context_schema, BaseModel) or is_dataclass(
+        context_schema,
+    )
+    if isinstance(context, dict) and schema_is_class:
+        return context_schema(**_filter_context_by_schema(context, graph))
+
+    return context
+
+
+_CACHE = {}
+
+
+def _filter_context_by_schema(context: dict[str, Any], graph: Pregel) -> dict[str, Any]:
+    if graph not in _CACHE:
+        _CACHE[graph] = graph.get_context_jsonschema()
+        if len(_CACHE) > 500:
+            _CACHE.popitem()
+    json_schema = _CACHE[graph]
+    if not json_schema or not context:
+        return context
+
+    # Extract valid properties from the schema
+    properties = json_schema.get("properties", {})
+    if not properties:
+        return context
+
+    # Filter context to only include parameters defined in the schema
+    filtered_context = {}
+    for key, value in context.items():
+        if key in properties:
+            filtered_context[key] = value
+
+    return filtered_context