feat(dynaflow): DynamoDB checkpointer for run/thread state (#280)

src/dynavec/__init__.py

@@ -23,6 +23,12 @@
 from .async_client import AsyncDynavec
 from .bm25 import BM25Index
 from .cache import BaseCache, DynamoDBCache, RedisCache, SemanticCache, warm_cache
+from .checkpoint import (
+    BaseCheckpointer,
+    Checkpoint,
+    DynamoDBCheckpointer,
+    MemoryCheckpointer,
+)
 from .client import Dynavec
 from .config import DynavecConfig
 from .credentials import AWSCredentials
@@ -178,6 +184,11 @@
     "get_run_logger",
     "set_default_run_logger",
     "log_run_event",
+    # Orchestration & Checkpointing (dynaflow)
+    "Checkpoint",
+    "BaseCheckpointer",
+    "MemoryCheckpointer",
+    "DynamoDBCheckpointer",
     # Ingestion
     "Record",
     "ingest",

src/dynavec/agents/__init__.py

@@ -2,6 +2,12 @@
 
 from __future__ import annotations
 
+from ..checkpoint import (
+    BaseCheckpointer,
+    Checkpoint,
+    DynamoDBCheckpointer,
+    MemoryCheckpointer,
+)
 from .base import (
     AgentResult,
     AgentStep,
@@ -35,4 +41,8 @@
     "mcp_tools_from_session",
     "mcp_tools_from_session_async",
     "tool",
+    "Checkpoint",
+    "BaseCheckpointer",
+    "MemoryCheckpointer",
+    "DynamoDBCheckpointer",
 ]

src/dynavec/agents/react.py

@@ -3,18 +3,69 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Sequence
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from ..chat.base import ChatModel, Message, Tool, ToolCall
 from .base import AgentResult, AgentStep, AgentTool
 from .registry import ToolRegistry
 
+if TYPE_CHECKING:
+    from ..checkpoint import BaseCheckpointer
+
 DEFAULT_REACT_SYSTEM_PROMPT = """You are a helpful and precise reasoning agent.
 You solve tasks step-by-step using a ReAct (Reason + Act) approach.
 When you need to look up information or perform calculations, invoke the appropriate tools.
 When you have the final answer to the user's goal, provide the clear and complete answer directly without calling any tools."""
 
 
+def _message_to_dict(msg: Message) -> dict[str, Any]:
+    return {
+        "role": msg.role,
+        "content": msg.content,
+        "tool_call_id": msg.tool_call_id,
+        "tool_calls": [
+            {"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in msg.tool_calls
+        ],
+    }
+
+
+def _dict_to_message(d: dict[str, Any]) -> Message:
+    tcs = [
+        ToolCall(id=t["id"], name=t["name"], arguments=t["arguments"])
+        for t in d.get("tool_calls", [])
+    ]
+    return Message(
+        role=d["role"],
+        content=d.get("content"),
+        tool_calls=tcs,
+        tool_call_id=d.get("tool_call_id"),
+    )
+
+
+def _step_to_dict(step: AgentStep) -> dict[str, Any]:
+    return {
+        "step_number": step.step_number,
+        "thought": step.thought,
+        "tool_calls": [
+            {"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in step.tool_calls
+        ],
+        "observations": list(step.observations),
+    }
+
+
+def _dict_to_step(d: dict[str, Any]) -> AgentStep:
+    tcs = [
+        ToolCall(id=t["id"], name=t["name"], arguments=t["arguments"])
+        for t in d.get("tool_calls", [])
+    ]
+    return AgentStep(
+        step_number=int(d["step_number"]),
+        thought=d.get("thought"),
+        tool_calls=tcs,
+        observations=list(d.get("observations") or []),
+    )
+
+
 class ReActAgent:
     """A ReAct-style agent that alternates between thinking, calling tools, and observing results."""
 
@@ -24,12 +75,14 @@ def __init__(
         tools: Sequence[AgentTool | Callable[..., Any]] | ToolRegistry | None = None,
         system_prompt: str | None = None,
         max_steps: int = 10,
+        checkpointer: BaseCheckpointer | None = None,
     ) -> None:
         self.model = model
         self.system_prompt = (
             system_prompt if system_prompt is not None else DEFAULT_REACT_SYSTEM_PROMPT
         )
         self.max_steps = max(1, max_steps)
+        self.checkpointer = checkpointer
 
         self._tools_map: dict[str, AgentTool] = {}
         self._chat_tools: list[Tool] = []
@@ -115,13 +168,47 @@ def _finalize_result(
             tool_calls_count=total_tool_calls,
         )
 
-    def run(self, goal: str, **kwargs: Any) -> AgentResult:
+    def run(
+        self,
+        goal: str,
+        *,
+        thread_id: str | None = None,
+        resume_from: str | None = None,
+        **kwargs: Any,
+    ) -> AgentResult:
         """Execute the ReAct loop synchronously until goal completion or max_steps."""
-        messages = self._init_messages(goal)
+        messages: list[Message] = []
         steps: list[AgentStep] = []
         total_tool_calls = 0
+        start_step = 1
+
+        if self.checkpointer is not None and thread_id is not None:
+            existing_cp = self.checkpointer.get(thread_id, resume_from)
+            if existing_cp is not None:
+                state = existing_cp.state
+                if "messages" in state:
+                    messages = [_dict_to_message(m) for m in state["messages"]]
+                if "steps" in state:
+                    steps = [_dict_to_step(s) for s in state["steps"]]
+                total_tool_calls = int(state.get("total_tool_calls", 0))
+                start_step = existing_cp.step + 1
 
-        for step_idx in range(1, self.max_steps + 1):
+        if not messages:
+            messages = self._init_messages(goal)
+            if self.checkpointer is not None and thread_id is not None:
+                self.checkpointer.put(
+                    thread_id,
+                    state={
+                        "goal": goal,
+                        "messages": [_message_to_dict(m) for m in messages],
+                        "steps": [],
+                        "total_tool_calls": 0,
+                    },
+                    node_id="init",
+                    step=0,
+                )
+
+        for step_idx in range(start_step, self.max_steps + 1):
             chat_res = self.model.invoke(
                 messages,
                 tools=self._chat_tools if self._chat_tools else None,
@@ -138,6 +225,20 @@ def run(self, goal: str, **kwargs: Any) -> AgentResult:
                     observations=[],
                 )
                 steps.append(step)
+                if self.checkpointer is not None and thread_id is not None:
+                    self.checkpointer.put(
+                        thread_id,
+                        state={
+                            "goal": goal,
+                            "messages": [_message_to_dict(m) for m in messages],
+                            "steps": [_step_to_dict(s) for s in steps],
+                            "total_tool_calls": total_tool_calls,
+                            "output": msg.content or "",
+                            "finished": True,
+                        },
+                        node_id=f"step_{step_idx}",
+                        step=step_idx,
+                    )
                 return AgentResult(
                     output=msg.content or "",
                     steps=steps,
@@ -158,15 +259,78 @@ def run(self, goal: str, **kwargs: Any) -> AgentResult:
             )
             steps.append(step)
 
-        return self._finalize_result(steps, total_tool_calls)
+            if self.checkpointer is not None and thread_id is not None:
+                self.checkpointer.put(
+                    thread_id,
+                    state={
+                        "goal": goal,
+                        "messages": [_message_to_dict(m) for m in messages],
+                        "steps": [_step_to_dict(s) for s in steps],
+                        "total_tool_calls": total_tool_calls,
+                        "finished": False,
+                    },
+                    node_id=f"step_{step_idx}",
+                    step=step_idx,
+                )
+
+        res = self._finalize_result(steps, total_tool_calls)
+        if self.checkpointer is not None and thread_id is not None:
+            self.checkpointer.put(
+                thread_id,
+                state={
+                    "goal": goal,
+                    "messages": [_message_to_dict(m) for m in messages],
+                    "steps": [_step_to_dict(s) for s in steps],
+                    "total_tool_calls": total_tool_calls,
+                    "output": res.output,
+                    "finished": False,
+                },
+                node_id="max_steps_reached",
+                step=self.max_steps,
+            )
+        return res
 
-    async def arun(self, goal: str, **kwargs: Any) -> AgentResult:
+    async def arun(
+        self,
+        goal: str,
+        *,
+        thread_id: str | None = None,
+        resume_from: str | None = None,
+        **kwargs: Any,
+    ) -> AgentResult:
         """Execute the ReAct loop asynchronously."""
-        messages = self._init_messages(goal)
+        messages: list[Message] = []
         steps: list[AgentStep] = []
         total_tool_calls = 0
+        start_step = 1
 
-        for step_idx in range(1, self.max_steps + 1):
+        if self.checkpointer is not None and thread_id is not None:
+            existing_cp = self.checkpointer.get(thread_id, resume_from)
+            if existing_cp is not None:
+                state = existing_cp.state
+                if "messages" in state:
+                    messages = [_dict_to_message(m) for m in state["messages"]]
+                if "steps" in state:
+                    steps = [_dict_to_step(s) for s in state["steps"]]
+                total_tool_calls = int(state.get("total_tool_calls", 0))
+                start_step = existing_cp.step + 1
+
+        if not messages:
+            messages = self._init_messages(goal)
+            if self.checkpointer is not None and thread_id is not None:
+                self.checkpointer.put(
+                    thread_id,
+                    state={
+                        "goal": goal,
+                        "messages": [_message_to_dict(m) for m in messages],
+                        "steps": [],
+                        "total_tool_calls": 0,
+                    },
+                    node_id="init",
+                    step=0,
+                )
+
+        for step_idx in range(start_step, self.max_steps + 1):
             chat_res = await self.model.ainvoke(
                 messages,
                 tools=self._chat_tools if self._chat_tools else None,
@@ -183,6 +347,20 @@ async def arun(self, goal: str, **kwargs: Any) -> AgentResult:
                     observations=[],
                 )
                 steps.append(step)
+                if self.checkpointer is not None and thread_id is not None:
+                    self.checkpointer.put(
+                        thread_id,
+                        state={
+                            "goal": goal,
+                            "messages": [_message_to_dict(m) for m in messages],
+                            "steps": [_step_to_dict(s) for s in steps],
+                            "total_tool_calls": total_tool_calls,
+                            "output": msg.content or "",
+                            "finished": True,
+                        },
+                        node_id=f"step_{step_idx}",
+                        step=step_idx,
+                    )
                 return AgentResult(
                     output=msg.content or "",
                     steps=steps,
@@ -205,4 +383,33 @@ async def arun(self, goal: str, **kwargs: Any) -> AgentResult:
             )
             steps.append(step)
 
-        return self._finalize_result(steps, total_tool_calls)
+            if self.checkpointer is not None and thread_id is not None:
+                self.checkpointer.put(
+                    thread_id,
+                    state={
+                        "goal": goal,
+                        "messages": [_message_to_dict(m) for m in messages],
+                        "steps": [_step_to_dict(s) for s in steps],
+                        "total_tool_calls": total_tool_calls,
+                        "finished": False,
+                    },
+                    node_id=f"step_{step_idx}",
+                    step=step_idx,
+                )
+
+        res = self._finalize_result(steps, total_tool_calls)
+        if self.checkpointer is not None and thread_id is not None:
+            self.checkpointer.put(
+                thread_id,
+                state={
+                    "goal": goal,
+                    "messages": [_message_to_dict(m) for m in messages],
+                    "steps": [_step_to_dict(s) for s in steps],
+                    "total_tool_calls": total_tool_calls,
+                    "output": res.output,
+                    "finished": False,
+                },
+                node_id="max_steps_reached",
+                step=self.max_steps,
+            )
+        return res