metaflow-durability 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- metaflow_durability-0.1.0/.github/workflows/release.yml +25 -0
- metaflow_durability-0.1.0/.github/workflows/test.yml +37 -0
- metaflow_durability-0.1.0/.gitignore +30 -0
- metaflow_durability-0.1.0/PKG-INFO +249 -0
- metaflow_durability-0.1.0/README.md +228 -0
- metaflow_durability-0.1.0/examples/coder_flow.py +180 -0
- metaflow_durability-0.1.0/examples/foreach_research.py +110 -0
- metaflow_durability-0.1.0/examples/handoff_agents.py +132 -0
- metaflow_durability-0.1.0/examples/parallel_agents.py +136 -0
- metaflow_durability-0.1.0/examples/research_flow.py +271 -0
- metaflow_durability-0.1.0/examples/retry_replay.py +82 -0
- metaflow_durability-0.1.0/examples/simple_agent.py +104 -0
- metaflow_durability-0.1.0/examples/triage_agents.py +129 -0
- metaflow_durability-0.1.0/pixi.lock +3111 -0
- metaflow_durability-0.1.0/pixi.toml +30 -0
- metaflow_durability-0.1.0/pyproject.toml +72 -0
- metaflow_durability-0.1.0/src/metaflow_durability/__init__.py +21 -0
- metaflow_durability-0.1.0/src/metaflow_durability/_version.py +24 -0
- metaflow_durability-0.1.0/src/metaflow_durability/agent_decorator.py +346 -0
- metaflow_durability-0.1.0/src/metaflow_durability/dispatch.py +155 -0
- metaflow_durability-0.1.0/src/metaflow_durability/durability.py +115 -0
- metaflow_durability-0.1.0/src/metaflow_durability/journal.py +177 -0
- metaflow_durability-0.1.0/src/metaflow_durability/journal_card.py +262 -0
- metaflow_durability-0.1.0/src/metaflow_durability/journal_decorator.py +158 -0
- metaflow_durability-0.1.0/src/metaflow_extensions/anaconda/config/mfextinit_anaconda.py +7 -0
- metaflow_durability-0.1.0/src/metaflow_extensions/anaconda/plugins/cards/__init__.py +183 -0
- metaflow_durability-0.1.0/src/metaflow_extensions/anaconda/plugins/mfextinit_anaconda.py +19 -0
- metaflow_durability-0.1.0/tests/durable_retry_flow.py +49 -0
- metaflow_durability-0.1.0/tests/test_durability.py +206 -0
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
name: release
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
release:
|
|
5
|
+
types: [published]
|
|
6
|
+
|
|
7
|
+
jobs:
|
|
8
|
+
test:
|
|
9
|
+
uses: ./.github/workflows/test.yml
|
|
10
|
+
|
|
11
|
+
publish:
|
|
12
|
+
needs: [test]
|
|
13
|
+
runs-on: ubuntu-latest
|
|
14
|
+
environment: pypi
|
|
15
|
+
permissions:
|
|
16
|
+
id-token: write # PyPI trusted publishing
|
|
17
|
+
steps:
|
|
18
|
+
- uses: actions/checkout@v5
|
|
19
|
+
with:
|
|
20
|
+
fetch-depth: 0
|
|
21
|
+
- uses: actions/setup-python@v6
|
|
22
|
+
with:
|
|
23
|
+
python-version: "3.13"
|
|
24
|
+
- run: pip install build && python -m build
|
|
25
|
+
- uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
name: test
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
pull_request:
|
|
5
|
+
push:
|
|
6
|
+
branches: [main]
|
|
7
|
+
merge_group:
|
|
8
|
+
workflow_call:
|
|
9
|
+
|
|
10
|
+
jobs:
|
|
11
|
+
test:
|
|
12
|
+
name: Test (Python ${{ matrix.python-version }})
|
|
13
|
+
runs-on: ubuntu-latest
|
|
14
|
+
strategy:
|
|
15
|
+
fail-fast: false
|
|
16
|
+
matrix:
|
|
17
|
+
python-version: ["3.11", "3.12", "3.13"]
|
|
18
|
+
steps:
|
|
19
|
+
- uses: actions/checkout@v5
|
|
20
|
+
with:
|
|
21
|
+
fetch-depth: 0
|
|
22
|
+
- uses: actions/setup-python@v6
|
|
23
|
+
with:
|
|
24
|
+
python-version: ${{ matrix.python-version }}
|
|
25
|
+
- run: pip install -e ".[pydantic-ai,dev]"
|
|
26
|
+
- run: pytest tests
|
|
27
|
+
|
|
28
|
+
# Wildcard job for branch protection: require only this check.
|
|
29
|
+
check:
|
|
30
|
+
name: Check all tests passed
|
|
31
|
+
if: always()
|
|
32
|
+
needs: [test]
|
|
33
|
+
runs-on: ubuntu-latest
|
|
34
|
+
steps:
|
|
35
|
+
- uses: re-actors/alls-green@release/v1
|
|
36
|
+
with:
|
|
37
|
+
jobs: ${{ toJSON(needs) }}
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
# Python
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*.egg-info/
|
|
5
|
+
dist/
|
|
6
|
+
build/
|
|
7
|
+
*.egg
|
|
8
|
+
|
|
9
|
+
# Virtual environments
|
|
10
|
+
.venv/
|
|
11
|
+
venv/
|
|
12
|
+
|
|
13
|
+
# Metaflow
|
|
14
|
+
.metaflow/
|
|
15
|
+
|
|
16
|
+
# Traces (session output, not source)
|
|
17
|
+
traces/
|
|
18
|
+
|
|
19
|
+
# OS
|
|
20
|
+
.DS_Store
|
|
21
|
+
|
|
22
|
+
# Editors / tools
|
|
23
|
+
.omo/
|
|
24
|
+
.kilo/
|
|
25
|
+
.vscode/
|
|
26
|
+
.idea/
|
|
27
|
+
|
|
28
|
+
src/metaflow_durability/_version.py
|
|
29
|
+
|
|
30
|
+
.env
|
|
@@ -0,0 +1,249 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: metaflow-durability
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Turns agent loops into durable, traceable, forkable work graphs on Metaflow
|
|
5
|
+
Project-URL: repository, https://github.com/anaconda-incubation/metaflow-durability
|
|
6
|
+
License: BSD-3-Clause
|
|
7
|
+
Requires-Python: >=3.11
|
|
8
|
+
Requires-Dist: litellm>=1.40
|
|
9
|
+
Requires-Dist: metaflow>=2.17.4
|
|
10
|
+
Requires-Dist: openai-function-calling>=2.0
|
|
11
|
+
Provides-Extra: dev
|
|
12
|
+
Requires-Dist: mypy; extra == 'dev'
|
|
13
|
+
Requires-Dist: pytest; extra == 'dev'
|
|
14
|
+
Requires-Dist: pytest-cov; extra == 'dev'
|
|
15
|
+
Requires-Dist: pytest-mock; extra == 'dev'
|
|
16
|
+
Requires-Dist: ruff; extra == 'dev'
|
|
17
|
+
Provides-Extra: pydantic-ai
|
|
18
|
+
Requires-Dist: markdown>=3.5; extra == 'pydantic-ai'
|
|
19
|
+
Requires-Dist: pydantic-ai-slim>=2.52; extra == 'pydantic-ai'
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# metaflow-durability
|
|
23
|
+
|
|
24
|
+
> [!WARNING]
|
|
25
|
+
> **This project is a proof of concept.** There is no guarantee that it will be maintained, and APIs may change or disappear without notice. You are welcome to discuss it in the `#incubation-team` Slack channel.
|
|
26
|
+
|
|
27
|
+
Agents in Metaflow steps, made durable. metaflow-durability has two independent parts; use either one, and neither needs the other:
|
|
28
|
+
|
|
29
|
+
- **`MetaflowDurability`**: a [pydantic-ai](https://ai.pydantic.dev) durable-execution capability. Your agent stays a plain pydantic-ai `Agent` that you call like a function inside any step. Every model call and tool call is journaled to the step's artifacts, and a retry replays the finished ones instead of paying for them again.
|
|
30
|
+
- **`@agent`**: a step decorator that runs its own LLM loop (via LiteLLM, an OpenAI-compatible client, or a co-stacked inference decorator) with plain Python tools. See [`@agent`](#agent-step-decorator) below.
|
|
31
|
+
|
|
32
|
+
## Quick start: MetaflowDurability
|
|
33
|
+
|
|
34
|
+
```
|
|
35
|
+
pip install "metaflow-durability[pydantic-ai]"
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
```python
|
|
39
|
+
from metaflow import FlowSpec, Parameter, retry, step
|
|
40
|
+
from pydantic import BaseModel
|
|
41
|
+
from pydantic_ai import Agent
|
|
42
|
+
|
|
43
|
+
from metaflow_durability import MetaflowDurability
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class Brief(BaseModel):
|
|
47
|
+
summary: str
|
|
48
|
+
sources: list[str]
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
researcher = Agent(
|
|
52
|
+
"anthropic:claude-sonnet-4-5",
|
|
53
|
+
name="researcher",
|
|
54
|
+
output_type=Brief,
|
|
55
|
+
tools=[web_search, fetch_page],
|
|
56
|
+
capabilities=[MetaflowDurability()],
|
|
57
|
+
)
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
class ResearchFlow(FlowSpec):
|
|
61
|
+
topic = Parameter("topic")
|
|
62
|
+
|
|
63
|
+
@retry(times=1)
|
|
64
|
+
@step
|
|
65
|
+
def start(self):
|
|
66
|
+
# a failed fetch replays finished calls, then runs the rest live
|
|
67
|
+
self.brief = researcher.run_sync(self.topic).output
|
|
68
|
+
self.next(self.end)
|
|
69
|
+
|
|
70
|
+
@step
|
|
71
|
+
def end(self):
|
|
72
|
+
print(self.brief.summary)
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
if __name__ == "__main__":
|
|
76
|
+
ResearchFlow()
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
If `fetch_page` fails partway through, `@retry` restarts the step. The searches, fetches, and model calls the first attempt finished are replayed from its journal, and only the rest run live. Nothing else in the flow changes: no extra decorator, and the agent code is ordinary pydantic-ai.
|
|
80
|
+
|
|
81
|
+
## How it works
|
|
82
|
+
|
|
83
|
+
- **The agent is a function call inside a step.** Code can run before and after it, a step can call several agents or the same agent several times, and plain code can route on the agent's typed output.
|
|
84
|
+
- **Every durable operation is journaled.** pydantic-ai's durable-execution builder hands each model request, tool call, tool discovery, and workspace call to `MetaflowDurability`. Each is recorded in the task's `agent_journal` artifact, keyed by agent name, which run of that agent it was in the task, and its position in that run.
|
|
85
|
+
- **Replay needs the name and the inputs to match.** On a retry, an operation is returned from the previous attempt's journal only if its name and a hash of its inputs (prompt, messages, settings, tool arguments) match. The first mismatch ends replay for that run and it continues live, so a changed prompt never gets a stale answer.
|
|
86
|
+
- **Nothing to attach.** The metaflow-durability extension adds a small step decorator to every step that hosts the journal and saves it as an artifact, on success or on failure. Outside Metaflow (tests, notebooks) the agent behaves as a plain agent.
|
|
87
|
+
|
|
88
|
+
## Requirements and constraints
|
|
89
|
+
|
|
90
|
+
- The agent needs a name: `Agent(name=...)`, or `MetaflowDurability(name=...)`, which takes precedence. The name keys the journal, so renaming it breaks replay of runs recorded under the old name.
|
|
91
|
+
- Two different `Agent` objects with the same name in one step raise an error. The same agent called again is fine.
|
|
92
|
+
- Tools run one at a time inside a durable run.
|
|
93
|
+
- A retry re-runs any tool call whose result was not yet recorded. Keep irreversible side effects (posting, publishing) in their own step.
|
|
94
|
+
- A hard kill (OOM, `SIGKILL`, preemption) loses that attempt's journal; exceptions, including `@timeout`, do not.
|
|
95
|
+
|
|
96
|
+
## Using a local or custom model endpoint
|
|
97
|
+
|
|
98
|
+
Any OpenAI-compatible endpoint works, including one served inside the task by an inference decorator such as `@llamacpp` or `@vllm`: wrap its client in a pydantic-ai model.
|
|
99
|
+
|
|
100
|
+
```python
|
|
101
|
+
from openai import AsyncOpenAI
|
|
102
|
+
from pydantic_ai.models.openai import OpenAIChatModel
|
|
103
|
+
from pydantic_ai.providers.openai import OpenAIProvider
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def local_model(base_url: str, model_name: str) -> OpenAIChatModel:
|
|
107
|
+
client = AsyncOpenAI(base_url=base_url, api_key="unused")
|
|
108
|
+
return OpenAIChatModel(model_name, provider=OpenAIProvider(openai_client=client))
|
|
109
|
+
```
|
|
110
|
+
|
|
111
|
+
There are two ways to hand that model to a durable agent:
|
|
112
|
+
|
|
113
|
+
**Build the agent in the step.** Use this when the endpoint only exists once the task is running. The name keeps the journal stable across attempts:
|
|
114
|
+
|
|
115
|
+
```python
|
|
116
|
+
@retry(times=1)
|
|
117
|
+
@step
|
|
118
|
+
def research(self):
|
|
119
|
+
researcher = Agent(
|
|
120
|
+
local_model(base_url, "granite-4.0-h-tiny"),
|
|
121
|
+
name="researcher",
|
|
122
|
+
capabilities=[MetaflowDurability()],
|
|
123
|
+
)
|
|
124
|
+
self.brief = researcher.run_sync(self.topic).output
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
**Register it, then pick it per call by name.** Use this when the endpoint is known when the agent is defined:
|
|
128
|
+
|
|
129
|
+
```python
|
|
130
|
+
researcher = Agent(
|
|
131
|
+
"anthropic:claude-sonnet-4-5",
|
|
132
|
+
name="researcher",
|
|
133
|
+
capabilities=[MetaflowDurability(models={"local": local_model(base_url, "granite")})],
|
|
134
|
+
)
|
|
135
|
+
|
|
136
|
+
researcher.run_sync(prompt, model="local") # registered model
|
|
137
|
+
researcher.run_sync(prompt) # the agent's default model
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
Passing an unregistered model instance at call time (`run_sync(model=...)` or `agent.override(model=...)`) raises an error inside a durable run, because pydantic-ai durable execution only accepts models it can find by name.
|
|
141
|
+
|
|
142
|
+
## Inspecting a run
|
|
143
|
+
|
|
144
|
+
The journal is an ordinary artifact, kept for every attempt:
|
|
145
|
+
|
|
146
|
+
```python
|
|
147
|
+
from metaflow import Task
|
|
148
|
+
|
|
149
|
+
task = Task("ResearchFlow/12/start/3", attempt=1)
|
|
150
|
+
for entry in task["agent_journal"].data["researcher"][0]:
|
|
151
|
+
source = "replayed" if entry["replayed"] else "live"
|
|
152
|
+
print(source, entry["name"], f"{entry['duration_ms']:.0f} ms")
|
|
153
|
+
```
|
|
154
|
+
|
|
155
|
+
Every step also gets an **Agent Journal** card:
|
|
156
|
+
|
|
157
|
+
- It updates live while the task runs: one row per operation, marked live or replayed, with duration, token counts, and a detail column.
|
|
158
|
+
- It ends with the agent's final result, rendered as markdown or highlighted JSON.
|
|
159
|
+
- A banner shows which attempt you are looking at: yellow while running, green for the final attempt, red for an attempt that failed. A failed attempt's card is re-rendered with its complete journal before the task exits.
|
|
160
|
+
- On a retried task, the final attempt's card is listed first (on S3, GCS, and Azure datastores).
|
|
161
|
+
|
|
162
|
+
```
|
|
163
|
+
python myflow.py card view <step_name>
|
|
164
|
+
```
|
|
165
|
+
|
|
166
|
+
## Examples
|
|
167
|
+
|
|
168
|
+
The examples use pydantic-ai's test model by default and a real model when `MODEL` is set (for example `bedrock:us.anthropic.claude-sonnet-4-6`).
|
|
169
|
+
|
|
170
|
+
- [`triage_agents.py`](examples/triage_agents.py): three agents chained in one step, with plain code routing on their typed output and the side effect in its own step.
|
|
171
|
+
- [`retry_replay.py`](examples/retry_replay.py): a tool fails on the first attempt; the `end` step prints which operations each attempt replayed and which ran live.
|
|
172
|
+
- [`foreach_research.py`](examples/foreach_research.py): the same agent across a foreach, then a synthesizer agent after the join.
|
|
173
|
+
- [`research_flow.py`](examples/research_flow.py): plain-code retrieval and citation checks around the agent, routing to a `deepen` step on weak briefs, a soft tool budget with pydantic-ai's `PrepareTools`, and a `publish` step that renders the report card. Needs `TAVILY_API_KEY`.
|
|
174
|
+
|
|
175
|
+
```
|
|
176
|
+
pixi run -e examples python examples/research_flow.py run
|
|
177
|
+
```
|
|
178
|
+
|
|
179
|
+
## `@agent` step decorator
|
|
180
|
+
|
|
181
|
+
`@agent` is the other half of metaflow-durability and does not use pydantic-ai. The decorator runs the whole LLM loop inside the step, before the step body: call the model, dispatch tool calls, append the results, and repeat until a final response or `max_iterations`.
|
|
182
|
+
|
|
183
|
+
```python
|
|
184
|
+
from metaflow import FlowSpec, Parameter, agent, step
|
|
185
|
+
|
|
186
|
+
|
|
187
|
+
class CityComparisonFlow(FlowSpec):
|
|
188
|
+
prompt = Parameter("prompt", default="Compare Tokyo and Paris by population density.")
|
|
189
|
+
|
|
190
|
+
@step
|
|
191
|
+
def start(self):
|
|
192
|
+
self.messages = [{"role": "user", "content": self.prompt}]
|
|
193
|
+
self.next(self.compare_cities)
|
|
194
|
+
|
|
195
|
+
@agent(
|
|
196
|
+
model="anthropic/claude-sonnet-4",
|
|
197
|
+
system_prompt="Call tools for data, then synthesize.",
|
|
198
|
+
tools=[city_stats, calculator],
|
|
199
|
+
max_iterations=10,
|
|
200
|
+
)
|
|
201
|
+
@step
|
|
202
|
+
def compare_cities(self):
|
|
203
|
+
self.next(self.end)
|
|
204
|
+
|
|
205
|
+
@step
|
|
206
|
+
def end(self):
|
|
207
|
+
print(self.final_response)
|
|
208
|
+
```
|
|
209
|
+
|
|
210
|
+
### Providing a model
|
|
211
|
+
|
|
212
|
+
Resolved in this order:
|
|
213
|
+
|
|
214
|
+
1. **An OpenAI-compatible client:** `@agent(client=OpenAI(base_url=..., api_key=...), tools=[...])`. Any object with `chat.completions.create()` works.
|
|
215
|
+
2. **A co-stacked inference decorator.** On [Anaconda AI Orchestration](https://www.anaconda.com/platform/ai-orchestration), stack `@llamacpp` or `@vllm` under `@agent` and leave out `model=`; the decorator finds the local OpenAI-compatible endpoint. No API keys, and it works air-gapped. Order: `@agent`, then the inference decorator, then `@resources`, closest to `@step`.
|
|
216
|
+
3. **A [LiteLLM](https://docs.litellm.ai/docs/providers) model string:** `@agent(model="bedrock/...", litellm_kwargs={"aws_region_name": "us-west-2"}, tools=[...])`. `litellm` is imported only on this path.
|
|
217
|
+
|
|
218
|
+
### Tools
|
|
219
|
+
|
|
220
|
+
Tools are plain functions with type hints and Google-style docstrings; schemas are generated with [openai-function-calling](https://github.com/jakecyr/openai-function-calling). The function name is the tool name, the docstring its description. Pass `tool_schemas=` to supply schemas yourself.
|
|
221
|
+
|
|
222
|
+
### Artifacts and card
|
|
223
|
+
|
|
224
|
+
| Artifact | Description |
|
|
225
|
+
|---|---|
|
|
226
|
+
| `messages` | Full conversation history |
|
|
227
|
+
| `iteration` | Number of LLM turns completed |
|
|
228
|
+
| `dispatched_tools` | Tool name, arguments, result, duration, and errors for every call |
|
|
229
|
+
| `final_response` | The final text, or `None` if `max_iterations` was reached |
|
|
230
|
+
|
|
231
|
+
Each `@agent` step also gets an agent card with summary stats, a tool dispatch table, and the final response. Stack your own `@card` with a different type to replace it.
|
|
232
|
+
|
|
233
|
+
Examples: [`simple_agent.py`](examples/simple_agent.py), [`parallel_agents.py`](examples/parallel_agents.py) (foreach), [`handoff_agents.py`](examples/handoff_agents.py) (sequential handoff).
|
|
234
|
+
|
|
235
|
+
## Setup
|
|
236
|
+
|
|
237
|
+
```
|
|
238
|
+
pip install metaflow-durability # @agent
|
|
239
|
+
pip install "metaflow-durability[pydantic-ai]" # MetaflowDurability
|
|
240
|
+
```
|
|
241
|
+
|
|
242
|
+
With [pixi](https://pixi.sh/): `pixi install`, then `pixi run test`. The `examples` environment adds Tavily and the pydantic-ai harness: `pixi run -e examples python examples/<name>.py run`. The examples load a `.env` if present (for example `MODEL`, `AWS_PROFILE`, `TAVILY_API_KEY`, and Metaflow service or datastore settings).
|
|
243
|
+
|
|
244
|
+
## Status
|
|
245
|
+
|
|
246
|
+
Early development.
|
|
247
|
+
|
|
248
|
+
- **Built:** `MetaflowDurability` recording and retry replay, and the Agent Journal card. `@agent` with tool dispatch, structured artifacts, and its card.
|
|
249
|
+
- **Planned:** resume and fork from a journal; keeping the journal through a hard kill; workspace and filesystem state.
|
|
@@ -0,0 +1,228 @@
|
|
|
1
|
+
# metaflow-durability
|
|
2
|
+
|
|
3
|
+
> [!WARNING]
|
|
4
|
+
> **This project is a proof of concept.** There is no guarantee that it will be maintained, and APIs may change or disappear without notice. You are welcome to discuss it in the `#incubation-team` Slack channel.
|
|
5
|
+
|
|
6
|
+
Agents in Metaflow steps, made durable. metaflow-durability has two independent parts; use either one, and neither needs the other:
|
|
7
|
+
|
|
8
|
+
- **`MetaflowDurability`**: a [pydantic-ai](https://ai.pydantic.dev) durable-execution capability. Your agent stays a plain pydantic-ai `Agent` that you call like a function inside any step. Every model call and tool call is journaled to the step's artifacts, and a retry replays the finished ones instead of paying for them again.
|
|
9
|
+
- **`@agent`**: a step decorator that runs its own LLM loop (via LiteLLM, an OpenAI-compatible client, or a co-stacked inference decorator) with plain Python tools. See [`@agent`](#agent-step-decorator) below.
|
|
10
|
+
|
|
11
|
+
## Quick start: MetaflowDurability
|
|
12
|
+
|
|
13
|
+
```
|
|
14
|
+
pip install "metaflow-durability[pydantic-ai]"
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
```python
|
|
18
|
+
from metaflow import FlowSpec, Parameter, retry, step
|
|
19
|
+
from pydantic import BaseModel
|
|
20
|
+
from pydantic_ai import Agent
|
|
21
|
+
|
|
22
|
+
from metaflow_durability import MetaflowDurability
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class Brief(BaseModel):
|
|
26
|
+
summary: str
|
|
27
|
+
sources: list[str]
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
researcher = Agent(
|
|
31
|
+
"anthropic:claude-sonnet-4-5",
|
|
32
|
+
name="researcher",
|
|
33
|
+
output_type=Brief,
|
|
34
|
+
tools=[web_search, fetch_page],
|
|
35
|
+
capabilities=[MetaflowDurability()],
|
|
36
|
+
)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
class ResearchFlow(FlowSpec):
|
|
40
|
+
topic = Parameter("topic")
|
|
41
|
+
|
|
42
|
+
@retry(times=1)
|
|
43
|
+
@step
|
|
44
|
+
def start(self):
|
|
45
|
+
# a failed fetch replays finished calls, then runs the rest live
|
|
46
|
+
self.brief = researcher.run_sync(self.topic).output
|
|
47
|
+
self.next(self.end)
|
|
48
|
+
|
|
49
|
+
@step
|
|
50
|
+
def end(self):
|
|
51
|
+
print(self.brief.summary)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
if __name__ == "__main__":
|
|
55
|
+
ResearchFlow()
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
If `fetch_page` fails partway through, `@retry` restarts the step. The searches, fetches, and model calls the first attempt finished are replayed from its journal, and only the rest run live. Nothing else in the flow changes: no extra decorator, and the agent code is ordinary pydantic-ai.
|
|
59
|
+
|
|
60
|
+
## How it works
|
|
61
|
+
|
|
62
|
+
- **The agent is a function call inside a step.** Code can run before and after it, a step can call several agents or the same agent several times, and plain code can route on the agent's typed output.
|
|
63
|
+
- **Every durable operation is journaled.** pydantic-ai's durable-execution builder hands each model request, tool call, tool discovery, and workspace call to `MetaflowDurability`. Each is recorded in the task's `agent_journal` artifact, keyed by agent name, which run of that agent it was in the task, and its position in that run.
|
|
64
|
+
- **Replay needs the name and the inputs to match.** On a retry, an operation is returned from the previous attempt's journal only if its name and a hash of its inputs (prompt, messages, settings, tool arguments) match. The first mismatch ends replay for that run and it continues live, so a changed prompt never gets a stale answer.
|
|
65
|
+
- **Nothing to attach.** The metaflow-durability extension adds a small step decorator to every step that hosts the journal and saves it as an artifact, on success or on failure. Outside Metaflow (tests, notebooks) the agent behaves as a plain agent.
|
|
66
|
+
|
|
67
|
+
## Requirements and constraints
|
|
68
|
+
|
|
69
|
+
- The agent needs a name: `Agent(name=...)`, or `MetaflowDurability(name=...)`, which takes precedence. The name keys the journal, so renaming it breaks replay of runs recorded under the old name.
|
|
70
|
+
- Two different `Agent` objects with the same name in one step raise an error. The same agent called again is fine.
|
|
71
|
+
- Tools run one at a time inside a durable run.
|
|
72
|
+
- A retry re-runs any tool call whose result was not yet recorded. Keep irreversible side effects (posting, publishing) in their own step.
|
|
73
|
+
- A hard kill (OOM, `SIGKILL`, preemption) loses that attempt's journal; exceptions, including `@timeout`, do not.
|
|
74
|
+
|
|
75
|
+
## Using a local or custom model endpoint
|
|
76
|
+
|
|
77
|
+
Any OpenAI-compatible endpoint works, including one served inside the task by an inference decorator such as `@llamacpp` or `@vllm`: wrap its client in a pydantic-ai model.
|
|
78
|
+
|
|
79
|
+
```python
|
|
80
|
+
from openai import AsyncOpenAI
|
|
81
|
+
from pydantic_ai.models.openai import OpenAIChatModel
|
|
82
|
+
from pydantic_ai.providers.openai import OpenAIProvider
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def local_model(base_url: str, model_name: str) -> OpenAIChatModel:
|
|
86
|
+
client = AsyncOpenAI(base_url=base_url, api_key="unused")
|
|
87
|
+
return OpenAIChatModel(model_name, provider=OpenAIProvider(openai_client=client))
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
There are two ways to hand that model to a durable agent:
|
|
91
|
+
|
|
92
|
+
**Build the agent in the step.** Use this when the endpoint only exists once the task is running. The name keeps the journal stable across attempts:
|
|
93
|
+
|
|
94
|
+
```python
|
|
95
|
+
@retry(times=1)
|
|
96
|
+
@step
|
|
97
|
+
def research(self):
|
|
98
|
+
researcher = Agent(
|
|
99
|
+
local_model(base_url, "granite-4.0-h-tiny"),
|
|
100
|
+
name="researcher",
|
|
101
|
+
capabilities=[MetaflowDurability()],
|
|
102
|
+
)
|
|
103
|
+
self.brief = researcher.run_sync(self.topic).output
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
**Register it, then pick it per call by name.** Use this when the endpoint is known when the agent is defined:
|
|
107
|
+
|
|
108
|
+
```python
|
|
109
|
+
researcher = Agent(
|
|
110
|
+
"anthropic:claude-sonnet-4-5",
|
|
111
|
+
name="researcher",
|
|
112
|
+
capabilities=[MetaflowDurability(models={"local": local_model(base_url, "granite")})],
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
researcher.run_sync(prompt, model="local") # registered model
|
|
116
|
+
researcher.run_sync(prompt) # the agent's default model
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
Passing an unregistered model instance at call time (`run_sync(model=...)` or `agent.override(model=...)`) raises an error inside a durable run, because pydantic-ai durable execution only accepts models it can find by name.
|
|
120
|
+
|
|
121
|
+
## Inspecting a run
|
|
122
|
+
|
|
123
|
+
The journal is an ordinary artifact, kept for every attempt:
|
|
124
|
+
|
|
125
|
+
```python
|
|
126
|
+
from metaflow import Task
|
|
127
|
+
|
|
128
|
+
task = Task("ResearchFlow/12/start/3", attempt=1)
|
|
129
|
+
for entry in task["agent_journal"].data["researcher"][0]:
|
|
130
|
+
source = "replayed" if entry["replayed"] else "live"
|
|
131
|
+
print(source, entry["name"], f"{entry['duration_ms']:.0f} ms")
|
|
132
|
+
```
|
|
133
|
+
|
|
134
|
+
Every step also gets an **Agent Journal** card:
|
|
135
|
+
|
|
136
|
+
- It updates live while the task runs: one row per operation, marked live or replayed, with duration, token counts, and a detail column.
|
|
137
|
+
- It ends with the agent's final result, rendered as markdown or highlighted JSON.
|
|
138
|
+
- A banner shows which attempt you are looking at: yellow while running, green for the final attempt, red for an attempt that failed. A failed attempt's card is re-rendered with its complete journal before the task exits.
|
|
139
|
+
- On a retried task, the final attempt's card is listed first (on S3, GCS, and Azure datastores).
|
|
140
|
+
|
|
141
|
+
```
|
|
142
|
+
python myflow.py card view <step_name>
|
|
143
|
+
```
|
|
144
|
+
|
|
145
|
+
## Examples
|
|
146
|
+
|
|
147
|
+
The examples use pydantic-ai's test model by default and a real model when `MODEL` is set (for example `bedrock:us.anthropic.claude-sonnet-4-6`).
|
|
148
|
+
|
|
149
|
+
- [`triage_agents.py`](examples/triage_agents.py): three agents chained in one step, with plain code routing on their typed output and the side effect in its own step.
|
|
150
|
+
- [`retry_replay.py`](examples/retry_replay.py): a tool fails on the first attempt; the `end` step prints which operations each attempt replayed and which ran live.
|
|
151
|
+
- [`foreach_research.py`](examples/foreach_research.py): the same agent across a foreach, then a synthesizer agent after the join.
|
|
152
|
+
- [`research_flow.py`](examples/research_flow.py): plain-code retrieval and citation checks around the agent, routing to a `deepen` step on weak briefs, a soft tool budget with pydantic-ai's `PrepareTools`, and a `publish` step that renders the report card. Needs `TAVILY_API_KEY`.
|
|
153
|
+
|
|
154
|
+
```
|
|
155
|
+
pixi run -e examples python examples/research_flow.py run
|
|
156
|
+
```
|
|
157
|
+
|
|
158
|
+
## `@agent` step decorator
|
|
159
|
+
|
|
160
|
+
`@agent` is the other half of metaflow-durability and does not use pydantic-ai. The decorator runs the whole LLM loop inside the step, before the step body: call the model, dispatch tool calls, append the results, and repeat until a final response or `max_iterations`.
|
|
161
|
+
|
|
162
|
+
```python
|
|
163
|
+
from metaflow import FlowSpec, Parameter, agent, step
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
class CityComparisonFlow(FlowSpec):
|
|
167
|
+
prompt = Parameter("prompt", default="Compare Tokyo and Paris by population density.")
|
|
168
|
+
|
|
169
|
+
@step
|
|
170
|
+
def start(self):
|
|
171
|
+
self.messages = [{"role": "user", "content": self.prompt}]
|
|
172
|
+
self.next(self.compare_cities)
|
|
173
|
+
|
|
174
|
+
@agent(
|
|
175
|
+
model="anthropic/claude-sonnet-4",
|
|
176
|
+
system_prompt="Call tools for data, then synthesize.",
|
|
177
|
+
tools=[city_stats, calculator],
|
|
178
|
+
max_iterations=10,
|
|
179
|
+
)
|
|
180
|
+
@step
|
|
181
|
+
def compare_cities(self):
|
|
182
|
+
self.next(self.end)
|
|
183
|
+
|
|
184
|
+
@step
|
|
185
|
+
def end(self):
|
|
186
|
+
print(self.final_response)
|
|
187
|
+
```
|
|
188
|
+
|
|
189
|
+
### Providing a model
|
|
190
|
+
|
|
191
|
+
Resolved in this order:
|
|
192
|
+
|
|
193
|
+
1. **An OpenAI-compatible client:** `@agent(client=OpenAI(base_url=..., api_key=...), tools=[...])`. Any object with `chat.completions.create()` works.
|
|
194
|
+
2. **A co-stacked inference decorator.** On [Anaconda AI Orchestration](https://www.anaconda.com/platform/ai-orchestration), stack `@llamacpp` or `@vllm` under `@agent` and leave out `model=`; the decorator finds the local OpenAI-compatible endpoint. No API keys, and it works air-gapped. Order: `@agent`, then the inference decorator, then `@resources`, closest to `@step`.
|
|
195
|
+
3. **A [LiteLLM](https://docs.litellm.ai/docs/providers) model string:** `@agent(model="bedrock/...", litellm_kwargs={"aws_region_name": "us-west-2"}, tools=[...])`. `litellm` is imported only on this path.
|
|
196
|
+
|
|
197
|
+
### Tools
|
|
198
|
+
|
|
199
|
+
Tools are plain functions with type hints and Google-style docstrings; schemas are generated with [openai-function-calling](https://github.com/jakecyr/openai-function-calling). The function name is the tool name, the docstring its description. Pass `tool_schemas=` to supply schemas yourself.
|
|
200
|
+
|
|
201
|
+
### Artifacts and card
|
|
202
|
+
|
|
203
|
+
| Artifact | Description |
|
|
204
|
+
|---|---|
|
|
205
|
+
| `messages` | Full conversation history |
|
|
206
|
+
| `iteration` | Number of LLM turns completed |
|
|
207
|
+
| `dispatched_tools` | Tool name, arguments, result, duration, and errors for every call |
|
|
208
|
+
| `final_response` | The final text, or `None` if `max_iterations` was reached |
|
|
209
|
+
|
|
210
|
+
Each `@agent` step also gets an agent card with summary stats, a tool dispatch table, and the final response. Stack your own `@card` with a different type to replace it.
|
|
211
|
+
|
|
212
|
+
Examples: [`simple_agent.py`](examples/simple_agent.py), [`parallel_agents.py`](examples/parallel_agents.py) (foreach), [`handoff_agents.py`](examples/handoff_agents.py) (sequential handoff).
|
|
213
|
+
|
|
214
|
+
## Setup
|
|
215
|
+
|
|
216
|
+
```
|
|
217
|
+
pip install metaflow-durability # @agent
|
|
218
|
+
pip install "metaflow-durability[pydantic-ai]" # MetaflowDurability
|
|
219
|
+
```
|
|
220
|
+
|
|
221
|
+
With [pixi](https://pixi.sh/): `pixi install`, then `pixi run test`. The `examples` environment adds Tavily and the pydantic-ai harness: `pixi run -e examples python examples/<name>.py run`. The examples load a `.env` if present (for example `MODEL`, `AWS_PROFILE`, `TAVILY_API_KEY`, and Metaflow service or datastore settings).
|
|
222
|
+
|
|
223
|
+
## Status
|
|
224
|
+
|
|
225
|
+
Early development.
|
|
226
|
+
|
|
227
|
+
- **Built:** `MetaflowDurability` recording and retry replay, and the Agent Journal card. `@agent` with tool dispatch, structured artifacts, and its card.
|
|
228
|
+
- **Planned:** resume and fork from a journal; keeping the journal through a hard kill; workspace and filesystem state.
|