dex-agents-langgraph 0.1.0__py3-none-any.whl

dex_agents_langgraph/__init__.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2026 Super Durable, Inc.
+#
+# Licensed under the Sustainable Use License 1.0.
+# You may not use this file except in compliance with the License.
+# See the LICENSE file in the repository root.
+#
+# SPDX-License-Identifier: LicenseRef-Sustainable-Use-1.0
+
+"""Durable Dex agent sessions for LangChain agents and LangGraph graphs.
+
+Build an agent with ``langchain.agents.create_agent`` and
+``DurableToolCallMiddleware``, or a custom ``StateGraph`` with
+``durable_tool_node``, and register it with ``langchain_agent_session_flow``
+or ``langgraph_agent_session_flow``. Each Flow execution is one agent session
+in Step mode: every model call is one ``RunModel`` Step, and every tool call
+runs in its own Step, in parallel, so a crash repeats at most the one call
+that was running. See the package README for the full guide.
+"""
+
+from dex_agents_langgraph.messages import (
+    framework_message_to_langchain_message,
+    framework_messages_to_langchain_messages,
+    langchain_message_to_framework_message,
+)
+from dex_agents_langgraph.runner import CompiledGraph, CompiledGraphSource, LangGraphAgentRunner
+from dex_agents_langgraph.session_flows import (
+    langchain_agent_session_flow,
+    langgraph_agent_session_flow,
+)
+from dex_agents_langgraph.tool_deferral import (
+    TOOL_CALL_IDEMPOTENCY_KEY_CONFIGURABLE,
+    DurableToolCallMiddleware,
+    DurableToolNode,
+    ToolCallExecutor,
+    awrap_durable_tool_call,
+    durable_tool_node,
+    get_tool_call_idempotency_key,
+    tool_effect_from_metadata,
+)
+
+__all__ = [
+    "TOOL_CALL_IDEMPOTENCY_KEY_CONFIGURABLE",
+    "CompiledGraph",
+    "CompiledGraphSource",
+    "DurableToolCallMiddleware",
+    "DurableToolNode",
+    "LangGraphAgentRunner",
+    "ToolCallExecutor",
+    "awrap_durable_tool_call",
+    "durable_tool_node",
+    "framework_message_to_langchain_message",
+    "framework_messages_to_langchain_messages",
+    "get_tool_call_idempotency_key",
+    "langchain_agent_session_flow",
+    "langchain_message_to_framework_message",
+    "langgraph_agent_session_flow",
+    "tool_effect_from_metadata",
+]

dex_agents_langgraph/_checkpoints.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2026 Super Durable, Inc.
+#
+# Licensed under the Sustainable Use License 1.0.
+# You may not use this file except in compliance with the License.
+# See the LICENSE file in the repository root.
+#
+# SPDX-License-Identifier: LicenseRef-Sustainable-Use-1.0
+
+"""The serialized LangGraph checkpoint that a turn stores while its tool calls run."""
+
+from __future__ import annotations
+
+import base64
+import json
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import ChannelVersions, Checkpoint, CheckpointMetadata
+from langgraph.checkpoint.memory import InMemorySaver
+
+# Bump when the JSON shape changes; restore must keep reading every earlier version.
+RUN_STATE_FORMAT = "dex-agents-langgraph.checkpoint.v1"
+
+
+# Checkpoint IDs follow each Worker's clock, so the largest ID may be stale; track the last put.
+class TurnCheckpointSaver(InMemorySaver):
+    def __init__(self, **saver_options: Any) -> None:
+        super().__init__(**saver_options)
+        self.latest_checkpoint_ids: dict[str, str] = {}
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        saved_config = super().put(config, checkpoint, metadata, new_versions)
+        self.latest_checkpoint_ids[config["configurable"]["checkpoint_ns"]] = checkpoint["id"]
+        return saved_config
+
+    def latest_checkpoint_config(self, thread_id: str, namespace: str = "") -> RunnableConfig:
+        return {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": namespace,
+                "checkpoint_id": self.latest_checkpoint_ids[namespace],
+            }
+        }
+
+
+@dataclass(frozen=True, kw_only=True)
+class PendingInterrupt:
+    interrupt_id: str
+    tool_call_ids: list[str]
+
+
+@dataclass(frozen=True, kw_only=True)
+class SuspendedGraphRunState:
+    thread_id: str
+    checkpoints: list[dict[str, Any]]
+    interrupts: list[PendingInterrupt]
+
+
+async def encode_suspended_graph_run_state(
+    saver: TurnCheckpointSaver,
+    thread_id: str,
+    interrupts: Sequence[PendingInterrupt],
+) -> str:
+    serde = saver.serde
+    checkpoints: list[dict[str, Any]] = []
+    for namespace in sorted(saver.latest_checkpoint_ids):
+        checkpoint_tuple = await saver.aget_tuple(
+            saver.latest_checkpoint_config(thread_id, namespace)
+        )
+        if checkpoint_tuple is None:
+            raise RuntimeError(f"checkpoint namespace {namespace!r} lost its latest checkpoint")
+        checkpoints.append(
+            {
+                "checkpoint_ns": namespace,
+                "checkpoint": _encode_typed(serde.dumps_typed(checkpoint_tuple.checkpoint)),
+                "metadata": _encode_typed(serde.dumps_typed(checkpoint_tuple.metadata)),
+                "pending_writes": [
+                    {
+                        "task_id": task_id,
+                        "channel": channel,
+                        "value": _encode_typed(serde.dumps_typed(value)),
+                    }
+                    for task_id, channel, value in checkpoint_tuple.pending_writes or ()
+                ],
+            }
+        )
+    return json.dumps(
+        {
+            "format": RUN_STATE_FORMAT,
+            "thread_id": thread_id,
+            "checkpoints": checkpoints,
+            "interrupts": [
+                {"interrupt_id": interrupt.interrupt_id, "tool_call_ids": interrupt.tool_call_ids}
+                for interrupt in interrupts
+            ],
+        },
+        separators=(",", ":"),
+    )
+
+
+def decode_suspended_graph_run_state(framework_run_state: str) -> SuspendedGraphRunState:
+    document = json.loads(framework_run_state)
+    if not isinstance(document, dict) or document.get("format") != RUN_STATE_FORMAT:
+        raise ValueError("the framework run state is not a LangGraph checkpoint of this adapter")
+    return SuspendedGraphRunState(
+        thread_id=str(document["thread_id"]),
+        checkpoints=list(document["checkpoints"]),
+        interrupts=[
+            PendingInterrupt(
+                interrupt_id=str(interrupt["interrupt_id"]),
+                tool_call_ids=[str(tool_call_id) for tool_call_id in interrupt["tool_call_ids"]],
+            )
+            for interrupt in document["interrupts"]
+        ],
+    )
+
+
+async def restore_suspended_graph_checkpoints(
+    saver: TurnCheckpointSaver,
+    run_state: SuspendedGraphRunState,
+) -> None:
+    serde = saver.serde
+    for stored in run_state.checkpoints:
+        checkpoint = serde.loads_typed(_decode_typed(stored["checkpoint"]))
+        metadata = serde.loads_typed(_decode_typed(stored["metadata"]))
+        namespace_config: RunnableConfig = {
+            "configurable": {
+                "thread_id": run_state.thread_id,
+                "checkpoint_ns": stored["checkpoint_ns"],
+            }
+        }
+        saved_config = await saver.aput(
+            namespace_config, checkpoint, metadata, checkpoint["channel_versions"]
+        )
+        # put_writes indexes writes by position, so each task's writes go in one call.
+        writes_by_task_id: dict[str, list[tuple[str, Any]]] = {}
+        for write in stored["pending_writes"]:
+            writes_by_task_id.setdefault(write["task_id"], []).append(
+                (write["channel"], serde.loads_typed(_decode_typed(write["value"])))
+            )
+        for task_id, writes in writes_by_task_id.items():
+            await saver.aput_writes(saved_config, writes, task_id)
+
+
+def _encode_typed(typed: tuple[str, bytes]) -> dict[str, str]:
+    return {"type": typed[0], "data": base64.b64encode(typed[1]).decode("ascii")}
+
+
+def _decode_typed(encoded: dict[str, str]) -> tuple[str, bytes]:
+    return encoded["type"], base64.b64decode(encoded["data"])

dex_agents_langgraph/messages.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2026 Super Durable, Inc.
+#
+# Licensed under the Sustainable Use License 1.0.
+# You may not use this file except in compliance with the License.
+# See the LICENSE file in the repository root.
+#
+# SPDX-License-Identifier: LicenseRef-Sustainable-Use-1.0
+
+"""Conversion between LangChain messages and stored agent messages."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+from typing import Any
+
+from dex_agents import AgentMessage, FrameworkMessage, MessageRole
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    ChatMessage,
+    HumanMessage,
+    ToolMessage,
+    message_to_dict,
+    messages_from_dict,
+)
+from pydantic_core import to_jsonable_python
+
+
+def langchain_message_to_framework_message(message: BaseMessage) -> FrameworkMessage:
+    """Convert one LangChain message into the item an agent session stores.
+
+    The runner calls it for every message a turn adds to the graph's message
+    state. ``framework_item`` holds ``langchain_core.messages.message_to_dict``
+    as JSON text, which ``framework_message_to_langchain_message`` reads back
+    unchanged on the next turn. The other fields are a display projection:
+
+    * ``HumanMessage`` and a ``ChatMessage`` with role ``"user"`` become
+      ``MessageRole.USER`` with ``text``.
+    * ``ToolMessage`` becomes ``MessageRole.TOOL`` with ``tool_call_id``,
+      ``tool_name``, and ``tool_output``.
+    * Every other message, including ``AIMessage`` and ``SystemMessage``,
+      becomes ``MessageRole.ASSISTANT``. An ``AIMessage`` with exactly one tool
+      call also sets ``tool_call_id``, ``tool_name``, and ``tool_arguments``.
+
+    Values that JSON cannot represent directly, such as a ``datetime`` in
+    ``response_metadata``, are stored in their JSON-compatible form, so they
+    come back as strings or plain containers.
+
+    Example:
+        >>> stored = langchain_message_to_framework_message(AIMessage("Done", id="ai-1"))
+        >>> stored.role, stored.text
+        (<MessageRole.ASSISTANT: 'assistant'>, 'Done')
+
+    Args:
+        message: A message from the graph's message state.
+
+    Returns:
+        The message with its JSON ``framework_item`` and display fields.
+
+    Raises:
+        TypeError: If a value in the message has no JSON-compatible form.
+    """
+    framework_item = json.dumps(
+        message_to_dict(message), ensure_ascii=False, default=to_jsonable_python
+    )
+    if isinstance(message, HumanMessage) or (
+        isinstance(message, ChatMessage) and message.role == "user"
+    ):
+        return FrameworkMessage(
+            role=MessageRole.USER, framework_item=framework_item, text=message.text
+        )
+    if isinstance(message, ToolMessage):
+        return FrameworkMessage(
+            role=MessageRole.TOOL,
+            framework_item=framework_item,
+            tool_call_id=message.tool_call_id,
+            tool_name=message.name,
+            tool_output=message.text,
+        )
+    text = message.text or None
+    if isinstance(message, AIMessage) and len(message.tool_calls) == 1:
+        tool_call = message.tool_calls[0]
+        return FrameworkMessage(
+            role=MessageRole.ASSISTANT,
+            framework_item=framework_item,
+            text=text,
+            tool_call_id=tool_call["id"],
+            tool_name=tool_call["name"],
+            tool_arguments=json.dumps(
+                tool_call["args"], ensure_ascii=False, default=to_jsonable_python
+            ),
+        )
+    return FrameworkMessage(role=MessageRole.ASSISTANT, framework_item=framework_item, text=text)
+
+
+def framework_message_to_langchain_message(message: AgentMessage) -> BaseMessage:
+    """Rebuild the LangChain message that a stored agent message holds.
+
+    The runner calls it for every stored message to rebuild the graph's input.
+    A message stored without a LangChain message ID receives
+    ``"dex-message-<sequence>"``, so later turns can tell stored messages from
+    new ones.
+
+    Example:
+        >>> history = [framework_message_to_langchain_message(m) for m in turn.messages]
+
+    Args:
+        message: A message from ``AgentTurnContext.messages``.
+
+    Returns:
+        The LangChain message that ``langchain_message_to_framework_message`` stored.
+
+    Raises:
+        ValueError: If ``framework_item`` is not a stored LangChain message.
+    """
+    item: Any = json.loads(message.framework_item)
+    if not isinstance(item, dict) or "type" not in item or "data" not in item:
+        raise ValueError(
+            f"agent message {message.sequence} does not hold a stored LangChain message"
+        )
+    restored = messages_from_dict([item])[0]
+    if restored.id is None:
+        restored.id = f"dex-message-{message.sequence}"
+    return restored
+
+
+def framework_messages_to_langchain_messages(
+    messages: Sequence[AgentMessage],
+) -> list[BaseMessage]:
+    """Rebuild every stored message of a session, in sequence order.
+
+    Args:
+        messages: ``AgentTurnContext.messages``.
+
+    Returns:
+        LangChain messages in the same order.
+
+    Raises:
+        ValueError: If a ``framework_item`` is not a stored LangChain message.
+    """
+    return [framework_message_to_langchain_message(message) for message in messages]