diff --git a/deepagents_plugin/README.md b/deepagents_plugin/README.md index 7f5d435b..d7afa4cf 100644 --- a/deepagents_plugin/README.md +++ b/deepagents_plugin/README.md @@ -26,6 +26,7 @@ one side. | [subagents](subagents) | Durability propagates across the agent tree — sub-agent model calls become activities with no per-sub-agent wiring. | | [streaming](streaming) | Stream model chunks to external subscribers via `streaming_topic` + `WorkflowStream`, keeping the durable result identical. | | [langsmith_tracing](langsmith_tracing) | Compose `DeepAgentsPlugin` with `LangSmithPlugin` for durable execution + LLM tracing. | +| [external_storage](external_storage) | Native External Storage for a 6 MiB tool result, plus an append-only S3 conversation transcript; scripted model, no API keys. | ## Prerequisites @@ -45,8 +46,10 @@ one side. > Temporal 1.34.0 or later and Python >= 3.11 — on older interpreters the group resolves to nothing > and the samples are skipped. -2. Configure a model provider. The samples use - `anthropic:claude-sonnet-4-5`, which needs an Anthropic API key: +2. Configure a model provider. Most samples use + `anthropic:claude-sonnet-4-5`, which needs an Anthropic API key. The + [external_storage](external_storage) sample uses a scripted model and needs + no provider credentials: ```bash export ANTHROPIC_API_KEY=... @@ -85,11 +88,13 @@ uv run deepagents_plugin/hello_world/run_worker.py uv run deepagents_plugin/hello_world/run_workflow.py ``` -The `langsmith_tracing` sample instead bundles the worker and starter into a -single driver: +The `langsmith_tracing` and `external_storage` samples instead bundle the worker +and starter into a single driver. The storage sample also needs a mock S3 service; +see their [setup instructions](../external_storage): ```bash uv run deepagents_plugin/langsmith_tracing/main.py +uv run deepagents_plugin/external_storage/main.py ``` ## Key Features Demonstrated @@ -109,6 +114,9 @@ uv run deepagents_plugin/langsmith_tracing/main.py - **Streaming** — forward model chunks to external subscribers while keeping the durable result unchanged. - **Observability** — compose with `LangSmithPlugin` for tracing. +- **Large payload storage** — configure native `ExternalStorage` on the client's + data converter, and retain conversation messages once in S3 with bounded model + context. ## Related diff --git a/deepagents_plugin/external_storage/README.md b/deepagents_plugin/external_storage/README.md new file mode 100644 index 00000000..7212c6eb --- /dev/null +++ b/deepagents_plugin/external_storage/README.md @@ -0,0 +1,85 @@ +# External Storage + +This sample uses S3 for two purposes: + +* A Deep Agent tool returns a **6 MiB** document. The tool result and the next + model activity input are stored in S3, with small references in workflow + history. +* The complete conversation is retained as **one immutable JSON object per + message**, including assistant tool calls and the full tool result. Each new + user turn passes a conversation ID, message offset, and the new question; + earlier turns are loaded inside the model activity. + +The workflow creates one agent and reuses it for every question. Each invocation +puts the conversation ID and message offset in the user message's +`additional_kwargs`. The plugin's `model_provider` returns the small worker-side +wrapper in `storage.py`, which reads that reference inside the model activity. +The agent loop stays in the workflow, and each model and tool call still runs as +its own plugin activity. No custom activity or Deep Agents backend is needed. + +Transcript objects live under `conversations//.json`. +Conditional S3 writes preserve existing messages and let a retried model call +reuse its committed reply, even after later turns have been added. One workflow +owns each conversation. System instructions remain model configuration, rather +than conversation messages. + +Only the latest **16 messages** are loaded as model context; older messages stay +in S3. This keeps both transcript storage and cumulative context reads linear in +the conversation's total bytes for this fixed window. Retaining every message +does not require sending every message to every model call. A production agent +can add token-based budgeting, summaries, or retrieval for older details. + +Temporal's separate `ExternalStorage` configuration uses its default 256 KiB +threshold. It offloads the large tool result and the following model activity +input. Ordinary conversation inputs stay small and inline, instead of saving a +new full-transcript snapshot on every turn. + +The S3 driver uses the local mock service from the +[general External Storage sample](../../external_storage). No AWS credentials, +Docker, or LLM API key is needed. The model replies are scripted; the real Deep +Agents loop, tool activity, model activities, and S3 transfers still run. + +## Run it + +Use Python **3.11 or later**: + +```bash +uv sync --python 3.13 --group deepagents --group external-storage +``` + +Start Temporal in one terminal: + +```bash +temporal server start-dev +``` + +Start the mock S3 service in another terminal. It listens on port 5000 and +creates the `temporal-payloads` bucket: + +```bash +uv run external_storage/s3.py +``` + +Run the demo: + +```bash +uv run deepagents_plugin/external_storage/main.py +``` + +The output includes the workflow ID, document size, transcript S3 prefix and +message count, and final answer. Inspect the workflow history with: + +```bash +temporal workflow show --workflow-id +``` + +Activity inputs and results above the threshold appear as native storage +references. To view externally stored payloads in the UI, run the +[storage-aware codec server](../../external_storage#5-optional-run-the-codec-server) +and set the UI's Remote Codec Endpoint to `http://localhost:8081`. + +For real S3, use an existing bucket and normal AWS credentials. Configure +compatible storage drivers on every client, worker, and replayer that reads the +history. Keep Temporal payload objects for as long as execution, reset, or replay +may need them, and keep transcript objects for the application's conversation +retention period. The local mock S3 service is in memory. diff --git a/deepagents_plugin/external_storage/client.py b/deepagents_plugin/external_storage/client.py new file mode 100644 index 00000000..ac54cdf6 --- /dev/null +++ b/deepagents_plugin/external_storage/client.py @@ -0,0 +1,56 @@ +"""Configure native S3 External Storage for the Deep Agents plugin.""" + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import replace + +import aioboto3 +from temporalio.client import Client +from temporalio.contrib.aws.s3driver import S3StorageDriver +from temporalio.contrib.aws.s3driver.aioboto3 import new_aioboto3_client +from temporalio.converter import DataConverter, ExternalStorage +from temporalio.deepagents import DeepAgentsPlugin +from temporalio.envconfig import ClientConfig + +from deepagents_plugin.external_storage.storage import ( + ConversationModel, + TranscriptStore, +) + +S3_ENDPOINT = "http://localhost:5000" +S3_BUCKET = "temporal-payloads" + + +@asynccontextmanager +async def connect_client( + *, + target_host: str | None = None, + s3_endpoint: str = S3_ENDPOINT, +) -> AsyncIterator[tuple[Client, TranscriptStore]]: + """Keep the S3 client alive while Temporal reads and writes payloads.""" + config = ClientConfig.load_client_connect_config() + config.setdefault("target_host", "localhost:7233") + if target_host is not None: + config["target_host"] = target_host + + session = aioboto3.Session() + async with session.client( + "s3", + endpoint_url=s3_endpoint, + aws_access_key_id="test", + aws_secret_access_key="test", + region_name="us-east-1", + ) as s3_client: + driver = S3StorageDriver( + client=new_aioboto3_client(s3_client), + bucket=S3_BUCKET, + ) + storage = ExternalStorage(drivers=[driver]) + data_converter = replace(DataConverter.default, external_storage=storage) + transcripts = TranscriptStore(s3_client, S3_BUCKET) + plugin = DeepAgentsPlugin( + data_converter=data_converter, + model_provider=lambda _: ConversationModel(store=transcripts), + ) + client = await Client.connect(**config, plugins=[plugin]) + yield client, transcripts diff --git a/deepagents_plugin/external_storage/main.py b/deepagents_plugin/external_storage/main.py new file mode 100644 index 00000000..c0702f93 --- /dev/null +++ b/deepagents_plugin/external_storage/main.py @@ -0,0 +1,49 @@ +"""Run an offline conversation with S3 transcripts and large tool payloads.""" + +import asyncio +import uuid + +from temporalio.worker import Worker + +from deepagents_plugin.external_storage.client import S3_BUCKET, connect_client +from deepagents_plugin.external_storage.workflow import ( + PAYLOAD_BYTES, + TASK_QUEUE, + ExternalStorageAgent, +) + + +async def main() -> None: + workflow_id = f"deepagents-external-storage-{uuid.uuid4()}" + questions = [ + "Read large-document.txt and acknowledge receiving it.", + "I am working on a project called Cedar.", + "The project uses Python.", + "It is due Friday.", + "Remind me of the project name, language, and deadline.", + ] + async with connect_client() as (client, _): + async with Worker( + client, + task_queue=TASK_QUEUE, + workflows=[ExternalStorageAgent], + ): + result = await client.execute_workflow( + ExternalStorageAgent.run, + questions, + id=workflow_id, + task_queue=TASK_QUEUE, + ) + + assert result.tool_result_bytes == PAYLOAD_BYTES + print(f"Workflow: {workflow_id}") + print(f"Tool result: {result.tool_result_bytes:,} bytes (6 MiB)") + print( + f"Transcript: s3://{S3_BUCKET}/conversations/" + f"{result.conversation_id}/ ({result.message_count} messages)" + ) + print(f"Final answer: {result.last_answer}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/deepagents_plugin/external_storage/storage.py b/deepagents_plugin/external_storage/storage.py new file mode 100644 index 00000000..f88f257e --- /dev/null +++ b/deepagents_plugin/external_storage/storage.py @@ -0,0 +1,171 @@ +"""Retain each conversation message once, outside Temporal's payload storage.""" + +import asyncio +import json +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +from botocore.exceptions import ClientError +from langchain_core.callbacks import ( + AsyncCallbackManagerForLLMRun, + CallbackManagerForLLMRun, +) +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import ( + AIMessage, + BaseMessage, + SystemMessage, + ToolMessage, + messages_from_dict, + messages_to_dict, +) +from langchain_core.outputs import ChatGeneration, ChatResult +from temporalio.deepagents.testing import FakeModel +from temporalio.exceptions import ApplicationError + +if TYPE_CHECKING: + from types_aiobotocore_s3.client import S3Client + +CONTEXT_MESSAGES = 16 + + +class TranscriptStore: + def __init__(self, client: "S3Client", bucket: str) -> None: + self.client = client + self.bucket = bucket + + async def get(self, conversation_id: str, index: int) -> BaseMessage | None: + try: + response = await self.client.get_object( + Bucket=self.bucket, + Key=f"conversations/{conversation_id}/{index:08d}.json", + ) + except ClientError as error: + if error.response["Error"]["Code"] == "NoSuchKey": + return None + raise + async with response["Body"] as body: + return messages_from_dict([json.loads(await body.read())])[0] + + async def put( + self, conversation_id: str, index: int, message: BaseMessage + ) -> BaseMessage: + try: + await self.client.put_object( + Bucket=self.bucket, + Key=f"conversations/{conversation_id}/{index:08d}.json", + Body=json.dumps(messages_to_dict([message])[0]).encode(), + ContentType="application/json", + IfNoneMatch="*", + ) + return message + except ClientError as error: + if error.response["Error"]["Code"] != "PreconditionFailed": + raise + existing = await self.get(conversation_id, index) + if existing is None: + raise RuntimeError( + "Transcript message disappeared after a conditional write" + ) + return existing + + async def read( + self, conversation_id: str, start: int, stop: int + ) -> list[BaseMessage]: + messages = await asyncio.gather( + *[self.get(conversation_id, index) for index in range(start, stop)] + ) + if any(message is None for message in messages): + raise ApplicationError("Transcript message is missing", non_retryable=True) + return [message for message in messages if message is not None] + + +class ConversationModel(BaseChatModel): + """Worker-side wrapper; the plugin still runs every model call as an activity.""" + + store: TranscriptStore + + @property + def _llm_type(self) -> str: + return "s3-conversation" + + def bind_tools(self, tools: Sequence[Any], **kwargs: Any) -> Any: + # The demo scripts its tool call rather than asking a real provider. + return self + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + raise NotImplementedError("Use the async model") + + async def _agenerate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: AsyncCallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + current = [m for m in messages if not isinstance(m, SystemMessage)] + reference = current[0].additional_kwargs + conversation_id = reference["conversation_id"] + offset = reference["message_offset"] + end = offset + len(current) + for index, message in enumerate(current, offset): + existing = await self.store.put(conversation_id, index, message) + # LangGraph can regenerate message IDs on replay; S3 positions are + # the stable identity. Tool call IDs remain part of the comparison. + if existing.model_dump(exclude={"id"}) != message.model_dump( + exclude={"id"} + ): + raise ApplicationError( + "Conflicting transcript message", non_retryable=True + ) + reply = await self.store.get(conversation_id, end) + if reply is None: + recent = await self.store.read( + conversation_id, max(0, end - CONTEXT_MESSAGES), end + ) + # Drop leading tool results if their calls fell outside the window. + while recent and isinstance(recent[0], ToolMessage): + recent.pop(0) + context = [m for m in messages if isinstance(m, SystemMessage)] + recent + response = await FakeModel([scripted_reply(recent)]).ainvoke(context) + # Conditional writes make committed replies reusable after retries. + reply = await self.store.put(conversation_id, end, response) + return ChatResult(generations=[ChatGeneration(message=reply)]) + + +def scripted_reply(messages: list[BaseMessage]) -> AIMessage: + latest = messages[-1] + if isinstance(latest, ToolMessage): + return AIMessage(content="I received the complete document.") + if str(latest.content).startswith("Read large-document.txt"): + return AIMessage( + content="", + tool_calls=[ + { + "name": "read_file", + "args": {"file_path": "large-document.txt"}, + "id": "call-read", + } + ], + ) + replies = { + "I am working on a project called Cedar.": "The project is called Cedar.", + "The project uses Python.": "Cedar uses Python.", + "It is due Friday.": "Cedar is due Friday.", + } + if str(latest.content).startswith("Remind me"): + context = "\n".join(str(message.content) for message in messages) + return AIMessage( + content=( + "Cedar is a Python project due Friday." + if all(fact in context for fact in ("Cedar", "Python", "Friday")) + else "Those details are outside my recent context." + ) + ) + return AIMessage(content=replies.get(str(latest.content), "Acknowledged.")) diff --git a/deepagents_plugin/external_storage/workflow.py b/deepagents_plugin/external_storage/workflow.py new file mode 100644 index 00000000..da32f6d0 --- /dev/null +++ b/deepagents_plugin/external_storage/workflow.py @@ -0,0 +1,93 @@ +"""Exercise native External Storage with a large tool result and a conversation.""" + +from dataclasses import dataclass +from datetime import timedelta + +from langchain_core.messages import HumanMessage, ToolMessage +from langchain_core.tools import tool +from temporalio import workflow +from temporalio.common import RetryPolicy +from temporalio.deepagents import create_temporal_deep_agent, tool_as_activity + +PAYLOAD_BYTES = 6 * 1024 * 1024 +TASK_QUEUE = "deepagents-external-storage" + + +@dataclass +class ExternalStorageResult: + tool_result_bytes: int + turns: int + last_answer: str + conversation_id: str + message_count: int + + +@tool +async def read_file(file_path: str) -> str: + """Read a mock document containing exactly 6 MiB of text.""" + del file_path + return "x" * PAYLOAD_BYTES + + +@workflow.defn +class ExternalStorageAgent: + @workflow.run + async def run(self, questions: list[str]) -> ExternalStorageResult: + reader = tool_as_activity( + read_file, + start_to_close_timeout=timedelta(seconds=30), + activity_options={"retry_policy": RetryPolicy(maximum_attempts=1)}, + ) + conversation_id = workflow.info().run_id + message_count = 0 + tool_result_bytes = 0 + last_answer = "" + agent = create_temporal_deep_agent( + model="fake:conversation", + tools=[reader], + system_prompt=( + "Answer using the recent conversation. For the first request, " + "read the document and acknowledge receiving it." + ), + activity_options={"start_to_close_timeout": timedelta(seconds=30)}, + ) + + for index, question in enumerate(questions): + result = await agent.ainvoke( + { + "messages": [ + HumanMessage( + content=question, + additional_kwargs={ + "conversation_id": conversation_id, + "message_offset": message_count, + }, + ) + ] + } + ) + if index == 0: + message = next( + message + for message in result["messages"] + if isinstance(message, ToolMessage) and message.name == "read_file" + ) + if not isinstance(message.content, str): + raise TypeError("Expected the document tool result to be text") + tool_result_bytes = len(message.content.encode()) + + answer = result["messages"][-1].content + if not isinstance(answer, str): + raise TypeError("Expected a text response from the model") + message_count += len(result["messages"]) + last_answer = answer + + if not questions: + raise ValueError("At least one question is required") + return ExternalStorageResult( + tool_result_bytes=tool_result_bytes, + turns=len(questions), + last_answer=last_answer, + conversation_id=conversation_id, + message_count=message_count, + ) diff --git a/tests/deepagents_plugin/external_storage_test.py b/tests/deepagents_plugin/external_storage_test.py new file mode 100644 index 00000000..b5922233 --- /dev/null +++ b/tests/deepagents_plugin/external_storage_test.py @@ -0,0 +1,196 @@ +"""Verify complete S3 transcripts without full-history activity inputs.""" + +import uuid +from collections.abc import Iterator +from datetime import timedelta + +import boto3 +import pytest +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage +from moto.server import ThreadedMotoServer +from temporalio.api.enums.v1 import EventType +from temporalio.api.sdk.v1 import ExternalStorageReference +from temporalio.client import Client +from temporalio.converter import DataConverter +from temporalio.deepagents import DeepAgentsPlugin +from temporalio.worker import Replayer, Worker + +from deepagents_plugin.external_storage import storage +from deepagents_plugin.external_storage.client import S3_BUCKET, connect_client +from deepagents_plugin.external_storage.storage import ( + CONTEXT_MESSAGES, + ConversationModel, +) +from deepagents_plugin.external_storage.workflow import ( + PAYLOAD_BYTES, + ExternalStorageAgent, +) +from tests.deepagents_plugin.helpers import ( + INVOKE_MODEL, + INVOKE_TOOL, + count_scheduled_activities, +) + + +@pytest.fixture +def s3_endpoint() -> Iterator[str]: + server = ThreadedMotoServer(ip_address="127.0.0.1", port=0, verbose=False) + server.start() + try: + host, port = server.get_host_and_port() + endpoint = f"http://{host}:{port}" + s3 = boto3.client( + "s3", + endpoint_url=endpoint, + aws_access_key_id="test", + aws_secret_access_key="test", + region_name="us-east-1", + ) + s3.create_bucket(Bucket=S3_BUCKET) + yield endpoint + finally: + server.stop() + + +@pytest.mark.parametrize("extra_turns", [0, 20]) +async def test_external_storage( + client: Client, s3_endpoint: str, extra_turns: int, monkeypatch: pytest.MonkeyPatch +) -> None: + queue = f"deepagents-storage-{uuid.uuid4()}" + target_host = client.service_client.config.target_host + questions = [ + "Read large-document.txt and acknowledge receiving it.", + "I am working on a project called Cedar.", + "The project uses Python.", + "It is due Friday.", + "Remind me of the project name, language, and deadline.", + *[f"Follow-up message {index}." for index in range(extra_turns)], + ] + contexts: list[list[BaseMessage]] = [] + scripted_reply = storage.scripted_reply + + def record_context(messages: list[BaseMessage]) -> AIMessage: + contexts.append(messages) + return scripted_reply(messages) + + monkeypatch.setattr(storage, "scripted_reply", record_context) + async with connect_client(target_host=target_host, s3_endpoint=s3_endpoint) as ( + stored_client, + transcripts, + ): + async with Worker( + stored_client, + task_queue=queue, + workflows=[ExternalStorageAgent], + max_cached_workflows=0, + ): + handle = await stored_client.start_workflow( + ExternalStorageAgent.run, + questions, + id=queue, + task_queue=queue, + execution_timeout=timedelta(seconds=90), + ) + result = await handle.result() + history = await handle.fetch_history() + + assert result.tool_result_bytes == PAYLOAD_BYTES + assert result.turns == len(questions) + assert result.last_answer == ( + "Acknowledged." if extra_turns else "Cedar is a Python project due Friday." + ) + assert result.message_count == 2 * len(questions) + 2 + counts = await count_scheduled_activities(handle) + assert counts[INVOKE_MODEL] == len(questions) + 1 + assert counts[INVOKE_TOOL] == 1 + assert all(len(context) <= CONTEXT_MESSAGES for context in contexts) + if extra_turns: + assert questions[0] not in [m.content for m in contexts[-1]] + + messages = await transcripts.read( + result.conversation_id, 0, result.message_count + ) + assert [m.content for m in messages if m.type == "human"] == questions + assert [m.type for m in messages[:4]] == ["human", "ai", "tool", "ai"] + assert isinstance(messages[1], AIMessage) + assert messages[1].tool_calls[0]["name"] == "read_file" + assert messages[2].content == "x" * PAYLOAD_BYTES + objects = await transcripts.client.list_objects_v2( + Bucket=S3_BUCKET, Prefix=f"conversations/{result.conversation_id}/" + ) + contents = objects["Contents"] + assert len(contents) == result.message_count + assert sum(obj["Size"] for obj in contents) < PAYLOAD_BYTES + 50_000 + + scheduled = { + event.event_id: event.activity_task_scheduled_event_attributes + for event in history.events + if event.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED + } + tool_result = next( + event.activity_task_completed_event_attributes.result.payloads[0] + for event in history.events + if event.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_COMPLETED + and scheduled[ + event.activity_task_completed_event_attributes.scheduled_event_id + ].activity_type.name + == INVOKE_TOOL + ) + model_inputs = [ + attrs.input.payloads[0] + for attrs in scheduled.values() + if attrs.activity_type.name == INVOKE_MODEL + ] + for payload in [tool_result, model_inputs[1]]: + reference = DataConverter.default.payload_converter.from_payload(payload) + assert isinstance(reference, ExternalStorageReference) + assert reference.driver_name == "aws.s3driver" + assert payload.ByteSize() < 1024 + for payload in model_inputs[2:]: + request = (await stored_client.data_converter.decode([payload]))[0] + assert len(request["messages"]) == 2 # System instructions + latest user. + assert not isinstance( + DataConverter.default.payload_converter.from_payload(payload), + ExternalStorageReference, + ) + assert "x" * 1024 not in history.to_json() + + async with connect_client(target_host=target_host, s3_endpoint=s3_endpoint) as ( + replay_client, + fresh_transcripts, + ): + assert ( + await fresh_transcripts.read( + result.conversation_id, 0, result.message_count + ) + == messages + ) + # Retry the first request after later turns: reuse the committed reply. + retry = await ConversationModel(store=fresh_transcripts).ainvoke([messages[0]]) + assert retry == messages[1] + assert len(contexts) == len(questions) + 1 + await Replayer( + workflows=[ExternalStorageAgent], + data_converter=replay_client.data_converter, + plugins=[DeepAgentsPlugin(data_converter=replay_client.data_converter)], + ).replay_workflow(history) + assert len(contexts) == len(questions) + 1 + + +async def test_transcript_messages_are_immutable( + s3_endpoint: str, client: Client +) -> None: + async with connect_client( + target_host=client.service_client.config.target_host, s3_endpoint=s3_endpoint + ) as (_, transcripts): + conversation_id = str(uuid.uuid4()) + original = HumanMessage(content="Keep this exact message.") + assert await transcripts.put(conversation_id, 0, original) == original + assert await transcripts.put(conversation_id, 0, original) == original + assert ( + await transcripts.put( + conversation_id, 0, HumanMessage(content="Different.") + ) + == original + ) + assert await transcripts.read(conversation_id, 0, 1) == [original]